Pytorch DDP多卡训练

本文最后更新于:2021年10月25日 晚上

Pytorch DDP多卡训练

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

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

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

1
2
3
4
5
6
7
8
9
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模型

1
2
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

1
2
3
4
5
6
7
8
9
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. 模型保存

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

5. 程序运行

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

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

1
2
3
4
5
6
7
8
9
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