torch分布式数据并行：torch.nn.parallel.DistributedDataParallel（DDP），代码书写步骤-Toy模板网

这篇具有很好参考价值的文章主要介绍了torch分布式数据并行：torch.nn.parallel.DistributedDataParallel（DDP），代码书写步骤。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

多进程做多卡训练；

1 初始化进程组：

2 当前进程所能用到的GPU卡的名称

3 将数据集随机分配到不同的GPU上

4 将train_sampler传入DataLoader中

5 将数据进行拷贝

6 模型放到GPU上

7 执行命令

8 模型保存

9 加载模型

10 注意事项

代码编写流程：

1 初始化进程组：

torch.distributed.init_process_group('nccl',worldsize=n_gpus,rank=args.local_rank)

‘nccl’ 指定GPU之间的通信方式；

world_size：当前这个节点上要用多少GPU卡；（当前节点就是当前机器）

rank: 当前进程在哪个GPU卡上，通过args.local_rank来获取，local_rank变量是通过外部指令传入的；（也可以通过环境变量来接收）

2 当前进程所能用到的GPU卡的名称

torch.cuda.set_device(args.local_rank)

上述指令作用相当于设置CUDA_VISBLE_DEVICES环境变量；

设置当前进程要用第几张卡；

3 将数据集随机分配到不同的GPU上

train_sampler = DistributedSampler(train_sampler)

该函数源码位于torch/utils/data/distributed.py

（返回一堆数据的索引，拿这些数据索引去dataloader中拿数据）

默认是Shuffle；

在每个epoch开始处，调用train_sampler.set_epoch(epoch)，使得数据充分打乱；

4 将train_sampler传入DataLoader中

train_dataloader = DataLoader(train_dataset,batch_size = BATCH_SIZE,collate_fn = collate_fn, sampler = train_sampler)

当DataLoader接受了一个自定义的sampler后，这就不用再传入shuffle=True了，这是互斥的。

5 将数据进行拷贝

data = data.cuda(args.local_rank)

当前节点上第几张卡

6 模型放到GPU上

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

7 执行命令

python -m torch.distributed.launch --nproc_per_node=n_gpus train.py

-m torch.distributed.launch 这个命令会构建多个进程，构建进程的数目通过 nproc_pre_node即每个设备的GPU数量来设置

launch这个程序会向每个train.py传入local_rank（0——n_gous-1）

8 模型保存

torch.save()在local_rank=0 这个GPU上保存，调用model.module.state_dict()来保存模型参数

例如

if step % save_step_interval == 0 and local_rank == 0:
    torch.save({
        'epoch': epoch_index,
        "step": step,
        "model_state_dict": model.module.state_dict(),
        "optimizer_state_dict": opyimizer.state_dict(),
        "loss": bce_loss
        },save_file)