训练流程简述
相比于DP,DDP能够启动多进程进行运算,从而大幅度提升计算资源的利用率。可以基于torch.distributed实现真正的分布式计算,具体的原理此处不再赘述。大致的流程如下:
初始化进程组。
创建分布式并行模型,每个进程都会有相同的模型和参数。
创建数据分发Sampler,使每个进程加载一个mini batch中不同部分的数据。
网络中相邻参数分桶,一般为神经网络模型中需要进行参数更新的每一层网络。
每个进程前向传播并各自计算梯度。
模型某一层的参数得到梯度后会马上进行通讯并进行梯度平均。
各GPU更新模型参数。
具体流程图如下:
图1
多机多卡数据并行训练
DistributedDataParallel进行多机多卡训练的优缺点
通信更快
:相比于DP,通信速度更快
负载相对均衡
:相比于DP,GPU负载相对更均衡
运行速度快
:因为通信时间更短,效率更高,能更快速的完成训练任务
代码改造点
引入多进程启动机制:初始化进程
引入几个变量:tcp协议,rank进程序号,worldsize开启的进程数量
分发数据:DataLoader中多了一个Sampler参数,避免不同进程数据重复
模型分发:DistributedDataParallel(model)
模型保存:在序号为0的进程下保存模型
import torch
class Net(torch.nn.Module):
model = Net().cuda()
### DistributedDataParallel Begin ###
model = torch.nn.parallel.DistributedDataParallel(Net().cuda())
### DistributedDataParallel End ###
分布式训练调测具体的代码适配操作过程和代码示例请参见
分布式调测适配及代码示例
章节。
文档还针对Resnet18在cifar10数据集上的分类任务,给出了分布式训练改造(DDP)的完整代码示例,供用户学习参考,具体请参见
分布式训练完整代码示例
。