PyTorch 使用DataParallel()实现多GPU训练

转载请注明作者和出处： http://blog.youkuaiyun.com/john_bh/

1.使用 DataParallel() 多卡训练

CLASS torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

其中包含三个主要的参数：module，device_ids和output_device。官方的解释如下：

• module即表示你定义的模型；
• device_ids表示你训练的device；
• output_device这个参数表示输出结果的device

当电脑或者服务器配备多卡的GPU，在上面跑模型的时候，当迭代次数或者epoch足够大的时候，通常会使用nn.DataParallel函数来用多个GPU来加速训练。在代码中加入以下代码，实现多卡训练：

device_ids = [0, 1]
net = torch.nn.DataParallel(net, device_ids=device_ids)

问题1：
加上这一行代码，在ternimal下执行watch -n 1 nvidia-smi后会发现确实会使用多个GPU来并行训练。但是第一块卡的显存会占用的更多一些，那么这是什么原因导致的？

torch.nn.DataParallel()最后一个参数output_device一般情况下是省略不写的，那么默认就是在device_ids[0]，也就是第一块卡上，进一步说也就是当你调用nn.DataParallel的时候，只是在你的input数据是并行的，但是你的output loss却不是这样的，每次都会在第一块GPU相加计算，这就造成了第一块GPU的负载远远大于剩余其他的显卡。

下面来具体讲讲nn.DataParallel中是怎么做的？
首先在前向过程中，你的输入数据会被划分成多个子部分（以下称为副本）送到不同的device中进行计算，而你的模型module是在每个device上进行复制一份，也就是说，输入的batch是会被平均分到每个device中去，但是你的模型module是要拷贝到每个