Pytorch 单机多卡训练DDP

多卡训练方式

1.DP——torch.nn.DataParallel
2.DDP——torch.nn.parallel.DistributedDataParallel 通俗一点讲就是用了4张卡训练，就会开启4个程序同时运行,每个卡运行总batch的1/4

方法比较

方法1简单，但是这种方式训练有不足之处。方法2要改动的地方比较多，但是速度更快。而且当模型很大的时候使用DataParallel我遇到了一个问题，报错说模型参数不在一个device上，这很有可能是单张卡放不下这些参数，但是具体的原因我也不清楚，改成DDP模式后即可正常运行。

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!

DDP使用

关键代码如下，根据自己需求修改。
代码解释
args.local_rank 指的是当前进程所用的设备，我这里用了4个卡，torch.distributed.launch运行时会更新4次local_rank
word_size 进程总数，一般一张卡开一个进程，可以理解为GPU个数，这里用到4张卡，开4个进程
train_sampler 需要把DataLoader里的shuffle=True换成train_sampler
dist.barrier() 同步屏障，在这里等待所有进程到此后才能进行训练
train_sampler.set_epoch(epoch) 保证每个进程拿到的数据都不同
–nproc_per_node=4 这个参数是给torch.distributed.launch的，让它开启4个进程，这个参数和word_size要一样

from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel
import torch.distributed as dist


parser.add_argument('--local_rank', type=int,default=0,help='local device i