使用Pytorch进行单机多卡分布式训练

一. torch.nn.DataParallel ?

pytorch单机多卡最简单的实现方法就是使用nn.DataParallel类,其几乎仅使用一行代码net = torch.nn.DataParallel(net)就可让模型同时在多张GPU上训练,它大致的工作过程如下图所示:
DataParallel前传与反传工作过程

每一个IterationForward过程中,nn.DataParallel都自动将输入按照gpu_batch进行split,然后复制模型参数到各个GPU上,分别进行前传后将得到网络输出,最后将结果concat到一起送往0号卡中。

Backward过程中,先由0号卡计算loss函数,通过loss.backward()得到损失函数相于各个gpu输出结果的梯度grad_l1 ... gradln,接下来0号卡将所有的grad_l送回对应的GPU中,然后GPU们分别进行backward得到各个GPU上面的模型参数梯度值gradm1 ... gradmn,最后所有参数的梯度汇总到GPU0卡进行update。

注:DataParallel的整个并行训练过程利用python多线程实现

由以上工作过程分析可知,nn.DataParallel有着这样几个无法避免的问题:

  1. 负载不均衡问题。gpu0所承担的任务明显要重于其他gpu
  2. 速度问题。每个iteration都需要复制模型且均从GPU0卡向其他GPU复制,通讯任务重且效率低;python多线程GIL锁导致的线程颠簸(thrashing)问题。
  3. 只能单机运行。由于单进程的约束导致。
  4. 只能切分batch到多GPU,而无法让一个model分布在多个GPU上。当一个模型过大,设置batchsize=1时其显存占用仍然大于单张显卡显存,此时就无法使用DataParallel类进行训练。

因此官方推荐使用torch.nn.DistributedDataParallel替代nn.DataParallel.

posted @ 2021-11-02 13:10  yinpengchen  阅读(1972)  评论(0编辑  收藏  举报