pytorch分布式训练方法总结

1 DP

1.1 原理

DP 基于单机多卡,所有设备都负责计算和训练网络,除此之外, device[0] (并非 GPU 真实标号而是输入参数 device_ids 首位) 还要负责整合梯度,更新参数。图 1 即为 GPU 0 作为 device[0] 的例子。从图中我们可以看出,有三个主要过程:

过程一(图中红色部分):各卡分别计算损失和梯度
过程二(图中蓝色部分):所有梯度整合到 device[0]
过程三(图中绿色部分):device[0] 进行参数更新,其他卡拉取 device[0] 的参数进行更新
所有卡都并行运算(图中红色),将梯度收集到 device[0](图中浅蓝色)和 device[0] 分享模型参数给其他 GPU(图中绿色)三个主要过程。
pytorch分布式训练方法总结

1.2 使用

DP的最大优点是使用非常简单, 只需要在原本单机版的代码中加上一行:

model = torch.nn.DataParallel(model)

使用DP时无须进行数据的划分, 数据会自动切分到不同的卡上, 每个卡上会分到batch_size/N 的数据, 程序中设定的batch_size为N个卡上的总和。

2 DDP

2.1 原理

3 Horovord

上一篇:检查 Tensorflow、Pytorch 是否可用 GPU


下一篇:实验室缺少GPU如何进行深度学习呢?精选云平台对比速速来看