Pytorch 并行训练（DP， DDP）的原理和应用

Pytorch 并行训练（DP， DDP）的原理和应用

单机多卡的情况下，torch.nn.DataParallel 的计算过程：

1、将 inputs 从主 GPU 分发到所有 GPU 上
2、将 model 从主 GPU 分发到所有 GPU 上
3、每个 GPU 分别独立进行前向传播，得到 outputs
4、将每个 GPU 的 outputs 发回主 GPU
5、在主 GPU 上，通过 loss function 计算出 loss，对 loss function 求导，求出损失梯度
6、计算得到的梯度分发到所有 GPU 上
7、反向传播计算参数梯度
8、将所有梯度回传到主 GPU，通过梯度更新模型权重
9、不断重复上面的过程

代码实现：

# API 如下

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

# 使用非常简单。一句代码就搞定。

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