PyTorch分布式训练方法
PyTorch多机多GPU并行训练方法及问题整理
参考MoCo训练源码：源码
torch.multiprocessing.spawn博客

单机多卡

初始化

torch.distributed.init_process_group()包含四个常用的参数
backend:后端，实际上是多个机器之间交换数据的协议
init_method:机器之间交换数据，需要指定一个主节点，而这个参数就是指定主节点的
world_size:介绍都是说是进程，实际就是机器的个数，例如两台机器一起训练的话，设置数值为2
rank:区分主节点和从节点，主节点为0，剩余的为了1-（N-1)，N为要使用的机器的数量，即 world_size

import torch.distributed as dist
dist.init_process_group(backend=dist_backend, init_method=args.dist_url,
                                world_size=world_size, rank=rank)

训练

from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel
dataset = your_dataset()
datasampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
dataloader = DataLoader(dataset, batch_size=batch_size_per_gpu, sampler=datasampler)
model = your_model()
model = DistributedDataPrallel(model, device_ids=[local_rank], output_device=local_rank)