使用pytorch分布式训练流程

引言

一般使用服务器进行多卡训练，这时候就需要使用pytorch的单机多卡的分布式训练方法，之前的api可能是

torch.nn.DataParallel

但是这个方法不支持使用多进程训练，所以一般使用下面的api来进行训练

torch.nn.parallel.DistributedDataParallel

这个api的执行效率会比上面的快很多，下面介绍代码的编写流程

参考资料：PyTorch多GPU分布式训练代码编写

代码编写流程

对于

torch.distributed.init_process_group(nccl", world_size=n_gpus, rank=args.local_rank)

初始化分布式系统，设置一些环境变量

nccl：各显卡之间的通讯方式，一般默认就行；

world_size: 使用显卡的数量；

rank：指定当前的进程的第几个gpu上，一般是训练的使用通过命令行输入。

model = DistributedDataParallel(model.cuda(args.locaL_rank), device_ids=[args.local_rank])

对原始模型进行包裹，先把model复制到当前的GPU上，在使用DDP进行包裹。

train_sampler = DistributedSampler(train_dataset)

对训练的数据来进行分配，随机加载到不同的GPU上。

这里要注意，因为是随机抽数据，但是当设置好随机种子，每次分配到不同的gpu上的数据其实是一样的，不利于训练，所以在每个周期开始处，调用train_sampler.set_epoch(epoch）可以使得数据充分打乱

在执行的时候，需要使用命令行来对一些环境设置进行传递，如下：

python -m torch.distributed.launch --nproc_per_node=n-n_gus train.py

模型的保存

注意事项：