pytorch 单机多卡训练distributedDataParallel

布满荆棘的人生 2022-09-03 06:15 205阅读 0赞

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