1.初始化进程组
dist.init_process_group(backend="nccl")
backend是后台利用nccl进行通信
2.使样本之间能够进行通信
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
3.创建ddp模型
model = DDP(model,device_ids=[local_rank],output_device=local_rank,find_unused_parameters=True)
获得local_rank(在运行launch时会传入一个local_rank参数)
local_rank = torch.distributed.local_rank()
torch.cuda.set_device(local_rank)
运行脚本
CUDA_VISIBLE_DEVICES=$device python -m torch.distributed.launch --nproc_per_node=1 train.py
保存模型(注意只需保存主进程上的模型，保存的是ddp模型的module)
if dist.get_rank() == 0: torch.save(model.module, "model.pth")