pytorch 多GPU训练 loss不能正常下降

在使用pytorch的DataParallel进行多GPU并行训练过程中，参考了很多博客，其中两个博客的内容有误， 会导致多GPU训练时loss不能正常下降，且运行过程中不会有报错， 下面做简单的说明。

有问题的博客链接：

https://blog.csdn.net/daniaokuye/article/details/79110365

https://blog.csdn.net/qq_19598705/article/details/80396325

以上两个博客的主要问题为将优化器也使用DataParallel包装：

optim_opt = config["optim"]
optimizer = optim.SGD(net.parameters(), lr = optim_opt["lr"], \
                momentum = optim_opt["momentum"], \
                nesterov = optim_opt["nesterov"], \
                weight_decay=optim_opt['weight_decay'])
optimizer = nn.DataParallel(optimizer, device_ids = [0, 1, 2, 3])

使用DataParallel包装后，optimizer的所有操作都需要获取module对象。 以上两个博客都忽略了优化器清空梯度时的操作：

optimizer.module.zero_grad()

如果使用：

optimizer.zero_grad()

程序不会报错， 但是训练时的loss不会正常下降。

事实上， 多卡训练过程中不必要将优化器也使用DataParallel包装， 可以采用如下的简单方式：

net = torch.nn.DataParallel(net.cuda(), device_ids = [i for i in range(GPU_NUM)])

后面就完全按照单卡方式就可以了。

需要注意的是在加载训练权重的时候， 多卡训练权重较单卡训练权重在关键字(key)中多了module字符串。

有关加载权重以及多卡训练负载均衡的问题， 可以参考：

简单的负载均衡方法

保存加载训练权重的一种方法