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(图中绿色)三个主要过程。
1.2 使用
DP的最大优点是使用非常简单, 只需要在原本单机版的代码中加上一行:
model = torch.nn.DataParallel(model)
使用DP时无须进行数据的划分, 数据会自动切分到不同的卡上, 每个卡上会分到batch_size/N 的数据, 程序中设定的batch_size为N个卡上的总和。