pytorch 单机多卡训练distributedDataParallel

pytorch单机多卡：从DataParallel到DistributedDataParallel
最近想做的实验比较多，于是稍微学习了一下和pytorch相关的加速方式。本人之前一直在使用DataParallel做数据并行，在今天浅浅的学了下apex之后，发现apex和DataParrallel并不兼容，由此开始了DistributedDataParallel的研究。至于在单机上DistributedDataParallel本身已经较DataParallel更优秀之类的内容，网上已经有较多详细的描述，这里就不再赘述。

单GPU无并行
我们要做的只有两件事：

告诉进程，可用的GPU的有些
把模型，损失函数，数据都放到GPU上
对于第一点，一般我采取两种不同的方式：

命令行在python前指定 CUDA_VISIBLE_DEVICES=0,1
在py文件内部，通过os.environ[‘CUDA_VISIBLE_DEVICES’] = ‘0.1’ 指定
第二点的模型和损失函数直接通过使用cuda()转化：

model = model.cuda()
Loss = Loss.cuda()

唯一需要注意的点在于，需要通过model.cuda()把model传到GPU上之后，再把model.parameters()传给optimizer

数据在train和validate的过程里的每个batch里转化

for data,label in dataloader:
    data = data.cuda()
    label = label.cuda()

DataParallel
DataParallel的使用比较简单，只需要在原本单GPU的基础上加上一句即可。

model = torch.nn.DataParallel(model)
model = model.cuda()

torch会在每个batch下，把数据分配给各个GPU进行计算。在这里，0号和1号GPU各分到batch-size/2的数据。每个GPU对分得的数据求完导数之后，会将导数传给主进程，主进程再进行参数的更新，然后把更新后的参数传给各个GPU。所以DataParallel多GPU的计算结果和单GPU的计算结果基本没有差别(每次更新用的数据个数等于batch-size).

DistributedDataParallel
DistributedDataParallel使用更复杂一些,这里只说单机的情况,涉及到多节点的一些参数和函数,只能说在我本人的机器上已经确认能够正常使用.

DistributedDataParallel的过程和DataParallel稍有不同,各子进程之间只进行导数的传播.具体的,以上面的设置为例,GPU0对分配给自己的数据求导,然后把该批数据的导数传给GPU1;GPU1对分配给自己的数据求导,然后把该批数据的导数传给GPU0,因此每个GPU都有该批次数据完整的导数,然后每个GPU均进行一次梯度下降,因为参数和梯度以及优化器都是一致的,所以每个GPU独立更新一步之后的参数仍然保持一致.

这里初看进行了重复的运算(每个GPU的参数更新是相同的),但是在DataParallel时,虽然只有一个主GPU进行参数的更新,但是其他GPU在此时只是在等待主GPU返回新的参数,所以这部分重复的计算不会导致运算时间的增加.

另一方面,DataParallel时传输数据用时为:各GPU传导数给主GPU,主GPU传更新后的参数给各GPU.而DistributedDataParallel只有GPU彼此间传导数这一步.

此时需要做的事情有:

一系列我没仔细研究的初始化设置
告诉每个进程使用GPU的index
在每个进程里把model和data放到对应的GPU上
为每个GPU生成对应的数据
初始化以及命令行参数
这一部分我不太了解,有兴趣的可以学习其他资料,比如pytorch 分布式训练 distributed parallel 笔记

我这里只描述一下我用了可以work的code.
在main()的开始,先初始化分布式通信:

torch.distributed.init_process_group(backend='nccl', init_method='env://')

通过命令行参数来获得进程index

parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int,default=0)
args = parser.parse_args()

如果原来的程序以及在使用argparse,则只需插入这里的第二行即可.
其中的参数’–local_rank’也不需要手动给出,会在下面通过命令行里辅助工具torch.distributed.launch自动给出:

python -m torch.distributed.launch --nproc_per_node=NUM_GPUS main.py [--arg1 --arg2 ...]

torch.distributed.launch将给每个进程分配一个编号’–local_rank’,下面我们将使用这个编号为各进程指定GPU.
NUM_GPUS为上面设置的可使用GPU的个数,这个例子里为2.

各进程GPU设置,model的设置
这里讲两种实现方式:

通过device=torch.device('cuda:{}'.format(args.local_rank)),来得到之前说的os.environ[‘CUDA_VISIBLE_DEVICES’] = ‘0.1’ 的第args.local_rank个GPU.在之后涉及到GPU时,使用tensor.to(device)即可.
通过torch.cuda.set_device(args.local_rank)来设置该进程使用的GPU序号,之后直接用tensor.cuda()即可
为每个GPU生成对应的数据
在DataParallel时,我们时通过一个完整的Dataloader来产生每个batch的数据,然后再自动把数据分配给各个GPU.使用DistributedDataParallel时,却是通过调用DistributedSampler来直接为各进程产生数据

train_sampler = torch.utils.data.distributed.DistributedSampler(train_data)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False,
                                               num_workers=2, pin_memory=True, sampler=train_sampler,)

这里需要设置shuffle=False,然后在每个epoch前,通过调用train_sampler.set_epoch(epoch)来达到shuffle的效果.注意事项
和DataParalle不同,DistributedDataParallel每个进程上的数据个数都等于batch-size,这导致每一次梯度下降实际用的数据量为batch-size*gpu-num.虽然一般情况下,batch越大计算越快,但是大batch不一定有最好的泛化,对于以及在单GPU或者DataParalle下以及对batch大小精调过的模型,设置batch为之前的GPU个数分之一才能复现之前结果
避免写入log文件发生冲突,只在某一个进程里用print,logger.可以通过args.local_rank==0来选择进程
计算loss或者accuracy时,要汇总各进程的信息,这里给出一个简单例子

    def reduce_tensor(tensor):
        # sum the tensor data across all machines
        dist.all_reduce(rt, op=dist.reduce_op.SUM)
        return rt

    output = model(input)
    loss = Loss(ouput, label) 
    accuracy = Accuracy(ouput, label)
    temp_tensor = torch.tensor([0])   # create a temp tensor to load batchsize which is scalar
    temp_tensor[0] = input.size(0)
    log_loss = reduce_tensor(loss.clone().detach_())
    log_acc = reduce_tensor(accuracy.clone().detach_()*batch_size)
    input_size = reduce_tensor(temp_tensor.clone().detach_())
    
    torch.cuda.synchronize()  # wait every process finish above transmission
    loss_total += log_loss.item()   
    if args.local_rank == 0:
        print('loss_total=',loss_total)
        print('accuracy=',log_acc/input_size)

因为刚写代码的时候对DistributedDataParallel的细节还不熟悉,所以按各GPU可能input.size()不一样在进行处理,通过了上面新建tensor来传输标量的笨办法来取得各GPU结果的权重.

原文链接：https://blog.csdn.net/weixin_39718268/article/details/105021631