多机多卡数据并行-DistributedDataParallel(DDP)

训练流程简述

相比于DP，DDP能够启动多进程进行运算，从而大幅度提升计算资源的利用率。可以基于torch.distributed实现真正的分布式计算，具体的原理此处不再赘述。大致的流程如下：

1. 初始化进程组。
2. 创建分布式并行模型，每个进程都会有相同的模型和参数。
3. 创建数据分发Sampler，使每个进程加载一个mini batch中不同部分的数据。
4. 网络中相邻参数分桶，一般为神经网络模型中需要进行参数更新的每一层网络。
5. 每个进程前向传播并各自计算梯度。
6. 模型某一层的参数得到梯度后会马上进行通讯并进行梯度平均。
7. 各GPU更新模型参数。

具体流程图如下：

图1 多机多卡数据并行训练

DistributedDataParallel进行多机多卡训练的优缺点

• 通信更快：相比于DP，通信速度更快
• 负载相对均衡：相比于DP，GPU负载相对更均衡
• 运行速度快：因为通信时间更短，效率更高，能更快速的完成训练任务

代码改造点

• 引入多进程启动机制：初始化进程
• 引入几个变量：tcp协议，rank进程序号，worldsize开启的进程数量
• 分发数据：DataLoader中多了一个Sampler参数，避免不同进程数据重复
• 模型分发：DistributedDataParallel(model)
• 模型保存：在序号为0的进程下保存模型

import torch
class Net(torch.nn.Module):
	pass

model = Net().cuda()

### DistributedDataParallel Begin ###
model = torch.nn.parallel.DistributedDataParallel(Net().cuda())
### DistributedDataParallel End ###

相关操作

• 分布式训练调测具体的代码适配操作过程和代码示例请参见分布式调测适配及代码示例章节。
• 文档还针对Resnet18在cifar10数据集上的分类任务，给出了分布式训练改造(DDP)的完整代码示例，供用户学习参考，具体请参见分布式训练完整代码示例。