Pytorch DataParallel and DistributedDataParallel

Pytorch DataParallel and DistributedDataParallel

最近试着使用Pytorch跑单机多卡训练，遇到了不少问题，做个总结和教程方便未来观看。我自己也是一个新手，很多东西总结的不好，有问题请多多指教，不懂的地方可以看参考文档，很多东西写的比我详细（本文只针对单机多卡训练，多机多卡训练未经过验证，请酌情观看）
环境：
python 3.7
pytorch 1.4.0

DataParallel

DataParallel是官方最早提供的一个库，使用非常简单，一行就够了

import torch.nn as nn

device = torch.device("cuda")
model = model.to(device)
model = nn.DataParallel(model)

DataParallel的缺点也很明显，只能在单台计算机上使用，而且速度很慢（相较于其他几种加速方式来说），最大的问题在于DataParallel采用了数据分布，但是loss是在某一块卡上计算的，这必然导致会出现比较严重的负载不均衡问题，其他卡用了很少的显存而主卡的显存使用很多，batch size稍微大一点主卡的显存就会炸了，因此更推荐使用其他的方式进行并行计算。

DistributedDataParallel

DistributedDataParallel（DDP）也是官方提供的一个并行计算库，相比于前者，该方式可以更好地进行多机多卡运算，更好的进行并行计算，负载均衡，运行效率也更高，支持的库和方法也更多，虽然使用起来较为麻烦，但对于追求性能来讲是一个更好的选择。

DDP会自动将库分配给n个进程并在n个gpu上运行，loss的计算也是在各个gpu上进行的，很大程度上解决了分配不均衡的问题。

模型设置

• 调用init设置通信后端
  用于指定通信后端，一般设置nccl就可以，这是nvidia的GPU的参数设定，更具体的设定请参考官方文档。

torch.distributed.init_process_group(backend='nccl')

• 添加local_rank参数
  local rank是为了DDP运行的时候进行节点进程制定，这个参数设定进去就可以，我们在自己进行计算时不需要人为给定参数，调用启动器的时候会自动设置

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1