pytorch指定gpu训练_DataParallel & DistributedDataParallel分布式训练

写在前面

今天跑测试代码的时候遇到了以下问题:对于同一个模型,相同的参数(同一个.pth文件),加不加

model = nn.DataParallel(model)

测试结果相差特别多(如下图所示),加了这句mIoU是第一个结果0.8847,没加是0.4929(以mIoU为例,可以看到其他各项指标也都掉的严重),所以决定花点心思把nn.DataParallel(model)搞清楚。

指标结果掉的严重

这里需要注意一下:多卡训练要考虑通信开销的, 是个trade-off的过程,不见得四块卡一定比两块卡快多少,训练到四块卡的时候可能io通信开销已经占了大头。

nn.DataParallel()

(我是单机多卡训练,首先送上官网nn.DataParallel()链接)

https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html?highlight=nn%20dataparallel#torch.nn.DataParallel​pytorch.org

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

• 参数定义

module是要放到多卡训练的模型;

device_ids数据类型是一个列表, 表示可用的gpu卡号;

output_devices数据类型也是列表,表示模型输出结果存放的卡号(如果不指定的话,默认放在0卡,这也是为什么多gpu训练并不是负载均衡的,一般0卡会占用的多,这里还涉及到一个小知识点——如果程序开始加os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3", 那么0卡(逻辑卡号)指的是2卡(物理卡号))。

• 如何多卡加速训练

讲原理：

网络在前向传播的时候会将model从主卡(默认是逻辑0卡)复制一份到所有的device上,input_data会在batch这个维度被分组后upload到不同的device上计算。在反向传播时，每个卡上的梯度会汇总到主卡上,求得梯度的均值后,再用反向传播更新单个GPU上的模型参数，最后将更新后的模型参数复制到剩余指定的GPU中进行下一轮的前向传播,以此来实现并行。有多少个gpu参与计算就有多少个gradient求平均,关于gradient求平均是否合理以及如何改进可以参考：

https://github.com/pytorch/pytorch/pull/7973/commits/c285b3626a7a4dcbbddfba1a6b217a64a3f3f3be​github.com

除了model和data,其他数据的分配规则如下:允许将Arbitrary positional和keyword添加到所有gpu上并行计算, 但有些数据类型是专门处理的。tensor将分配到指定的卡上,默认为0(这个地方我可能理解的不对,原话是“tensors will be scatteredon dim specified (default 0)”如果有人有其他正确的理解,望不吝赐教)。对元组、列表和字典类型的数据浅拷贝后复