PyTorch 通过两种方式可以进行多GPU训练: DataParallel, DistributedDataParallel. 当使用DataParallel的时候, 梯度的计算结果和在单卡上跑是一样的, 对每个数据计算出来的梯度进行累加. 当使用DistributedDataParallel的时候, 每个卡单独计算梯度, 然后多卡的梯度再进行平均.
下面是实验验证:
DataParallel
import torch
import os
import torch.nn as nn
def main():
model = nn.Linear(2, 3).cuda()
model = torch.nn.DataParallel(model, device_ids=[0, 1])
input = torch.rand(2, 2)
labels = torch.tensor([[1, 0, 0], [0, 1, 0]]).cuda()
(model(input) * labels).sum().backward()
print('input', input)
print([p.grad for p in model.parameters()])
if __name__=="__main__":
main()
执行CUDA_VISIBLE_DEVICES=0,1 python t.py
可以看到输出, 代码中对两个样本分别求梯度, 梯度等于样本的值, DataParallel把两个样本的梯度累加起来在不同GPU中同步.
input tensor([[0.4362, 0.4574],
[0.2052, 0.2362]])
[tensor([[0.4363, 0.4573],
[0.2052, 0.2362],