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字符串。

有關加載權重以及多卡訓練負載均衡的問題， 可以參考：

簡單的負載均衡方法

保存加載訓練權重的一種方法