pytorch 分布式訓練

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")