pytorch多GPU训练的两种方式了解 nn.DataParallel，distributedDataparallel

真的没想到随手写的一篇小笔记会引起那么多关注，真是瑟瑟发抖。

读研之后，才开始接触pytorch, 理解的难免有很多错误，感谢各位大佬指出的错误和提出的意见，我会慢慢修改的。评论区有大佬说nvidia的apex.distributeddataparallel也是一个很好的选择，

https://github.com/NVIDIA/apex/blob/master/examples/imagenet/main_amp.py​ github.com

这里是文档里提供的一个简单例子，感兴趣的可以参考一下。

----------------------------分割线-----------------------------------

之前单机跑模型的时候，发现真真真的。。。。。是太慢了！像cityscape这种370g的数据集，等一个模型跑完，估计都该过年了！！！然后自然而然就想起来去用Dataparallel函数加速并行。

先来看一下如何使用这个函数，

if torch.cuda.device_count() > 1:#判断是不是有多个GPU
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    # 就这一行
    model = nn.DataParallel(model,device_ids=range(torch.cuda.device_count()))

看起来非常简单对不对，就一行代码就解决了，但是果真如此吗？在我使用的时候发现，GPU会出现负载不均衡的问题。什么意思呢?就是第一个GPU（12GB）可能占用了10GB，剩余的GPU却只使用了2GB。这就是极大的浪费啊！！！家里没矿并且想尽快得到结果的我就想解决这个问题，经过查阅资料发现，这个问题的原因是，当你在数据并行的时候，你的loss却不是这样的，每次都会在第一个GPU相加计算，这就造成了第一个GPU的负载远远大于剩余其他的显卡。

那么如何完美的解决这个问题呢？

我找到了两种办法，不妨一试：

1.pytorch-encoding

https://github.com/zhanghang1989/PyTorch-Encoding​ github.com

这个library添加了loss并行计算的功能，当然除此之外，还有很多别的操作（比如 Synchronize BN），但是和这次的主题没有关系，所以就不多赘述。这个库写的非常好，感兴趣的可以看看。

from utils.encoding importDataParallelModel, DataParallelCriterion
model =DataParallelModel(model)
criterion =DataParallelCriterion(criterion)

参考例子是这个，

https://github.com/speedinghzl/pytorch-segmentation-toolbox/blob/master/train.py​ github.com

2. distributedDataparallel

Distributed communication package - torch.distributed​ pytorch.org

This is the highly recommended way to use DistributedDataParallel, with multiple processes, each of which operates on a single GPU. This is currently the fastest approach to do data parallel training using PyTorch and applies to both single-node(multi-GPU) and multi-node data parallel training. It is proven to be significantly faster than torch.nn.DataParallelfor single-node multi-GPU data parallel training.

第二种方法就是今天要讲的主题啦，这个函数的主要目的是为了多机多卡加速的，但是单机多卡也是没问题的。相比于之前的Dataparallel，新的函数更加优雅，速度更加快（这一点官方文档里有提到），而且不会出现负载不均衡的问题，唯一的小缺点可能就是配置稍微有点小麻烦。

下面主要介绍一下我的调用过程，我粗略的把他分成三步：

1.初始化

#初始化使用nccl后端（这个），当然还有别的后端，可以查看官方文档，介绍的比较清晰
torch.distributed.init_process_group(backend="nccl")

2.使用DistributedSampler

起先我是不太懂为什么要修改dataloader的，只是默默跟着改就完事了。后来看了别的文章才知道，DDP并不会自动shard数据

1. 如果自己写数据流，得根据torch.distributed.get_rank()去shard数据，获取自己应用的一份
2. 如果用Dataset API，则需要在定义Dataloader的时候用DistributedSampler 去shard：

train_dataset = self.dataset(
 self.opt.data_path, train_filenames, self.opt.height, self.opt.width,
 self.opt.frame_ids, 4, is_train=True, img_ext=img_ext)

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)

train_loader = DataLoader(
train_dataset, self.opt.batch_size, False,
num_workers=self.opt.num_workers, pin_memory=True, drop_last=True,sampler=train_sampler)

另外多说一点，在使用DataLoader时，别忘了设置pip_memory=true，为什么呢？且看下面的解释，

多GPU训练的时候注意机器的内存是否足够(一般为使用显卡显存x2)，如果不够，建议关闭pin_memory(锁页内存)选项。
采用 DistributedDataParallel多GPUs训练的方式比 DataParallel更快一些，如果你的Pytorch编译时有nccl的支持，那么最好使用DistributedDataParallel方式。
关于什么是锁页内存：
pin_memory就是锁页内存，创建DataLoader时，设置pin_memory=True，则意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存就会更快一些。
主机中的内存，有两种存在方式，一是锁页，二是不锁页，锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换（注：虚拟内存就是硬盘），而不锁页内存在主机内存不足时，数据会存放在虚拟内存中。显卡中的显存全部是锁页内存,当计算机的内存充足的时候，可以设置pin_memory=True。当系统卡住，或者交换内存使用过多的时候，设置pin_memory=False。因为pin_memory与电脑硬件性能有关，pytorch开发者不能确保每一个炼丹玩家都有高端设备，因此pin_memory默认为False。

总结就一个字，快！！！

3.分布式训练

model=torch.nn.parallel.DistributedDataParallel(model)

到这里你以为就结束了吗？？其实并没有，还剩下一点工作需要做，就是在命令行里加上这样一句话，

python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE  yourscript.py

但是这样做完之后可能还会报错，然后我就找到了torch.distributed.launch这个启动文件，发现我错过了一条很重要的提示，如果需要用这个函数的话，推荐大家先看一下 distributed.launch里面的内容，然后去看官方的文档解释。

In your training program, you must parse the command-line argument: --local_rank=LOCAL_PROCESS_RANK, which will be provided by this module.If your training program uses GPUs, you should ensure that your code only runs on the GPU device of LOCAL_PROCESS_RANK. This can be done by: Parsing the local_rank argument

就是在外部的命令中制定–local_rank(int),然后就可以顺利的跑下去了。到此为止，这个问题就算解决了，如果后续再碰到问题，我会继续补充的。

参考文章

https://zhuanlan.zhihu.com/p/43424629

https://www.cnblogs.com/jfdwd/p/11196439.html

https://cloud.tencent.com/developer/article/1512508