Pytorch DDP多卡训练

DDP指的是DistributedDataParallel，即支持多机多卡的Pytorch官方库，位于torch.distributed中。

本文简单介绍使用DDP进行单机多卡训练的步骤。

1. 开头导入库、获取local_rank、初始化backend

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# 有的用parser获取LOCAL_RANK的已经过时了
local_rank = int(os.environ["LOCAL_RANK"]) 
torch.cuda.set_device(local_rank)
dist.init_process_group(backend='nccl')  # 一般使用的后端为nccl
device = torch.device("cuda", local_rank)

2. 初始化DDP模型

model = YourModel().to(device)  # 要在把模型放进GPU之后再DDP
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

3. 数据处理

使用DDP的batchsize是单张卡上的batchsize，假如你用一张卡训练的batchsize是64，然后想分散在四张卡上，那么batchsize应该设置成16

train_iter = torchvision.datasets.CIFAR10(root='./data', train=True)
train_sampler = torch.utils.data.distributed.DistributedSampler(train_iter)
train_loader = torch.utils.data.DataLoader(train_iter, batch_size=batch_size, sampler=train_sampler)  # 注意sampler！

# 在epoch中
for epoch in range(xx):
    trainloader.sampler.set_epoch(epoch)  # 注意set_epoch
    for xxx in train_loader:
        pass

4. 模型保存

if dist.get_rank() == 0:  # 先判断是主卡，免得所有卡都保存一次
    torch.save(model.module, "saved_model.pth")  # 保存的是DDP.module

5. 程序运行

# nproc_per_node就是使用多少张卡
# 可使用来选择多卡机器上的某几张卡CUDA_VISIBLE_DEVICES=4,5,6,7
python -m torch.distributed.launch --nproc_per_node 8 main.py

*6. 使用Pytorch推荐的spawn方式

def demo_fn(rank, world_size):  # 两个参数，第一个参数一定是rank
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    pass

import torch.multiprocessing as mp
mp.spawn(demo_fn,
         args=(world_size,), # 传参给demo_fn，调用demo_fn的时候还会有一个参数表示rank
         nprocs=world_size,  # 用多少卡
         join=True)

参考文献

[原创][深度][PyTorch] DDP系列第一篇：入门教程 - 知乎 (zhihu.com)

Multiprocessing package - torch.multiprocessing — PyTorch 1.10.0 documentation