torch.distributed使用总结 单机多卡并行训练demo

转载：https://zhuanlan.zhihu.com/p/468784034

一、代码总览

一段完整的代码以及程序启动命令

训练代码

import os
import argparse
import torch
from torch.nn import SyncBatchNorm
from torch.nn.parallel import DistributedDataParallel
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler

parser = argparse.ArgumentParser(description='training')
parser.add_argument('--local_rank', type=int, help='local rank for dist')
args = parser.parse_args()

print(os.environ['MASTER_ADDR'])
print(os.environ['MASTER_PORT'])
world_size = torch.cuda.device_count()
local_rank = args.local_rank

dist.init_process_group(backend='nccl')

torch.cuda.set_device(local_rank)

train_dataset = Dataset(...)
train_sampler = DistributedSampler(train_dataset)
train_loader = Dataloader(dataset=train_dataet, sampler=train_sampler, shuffle=False)

val_set = Dataset()
val_loader = Dataloader(dataset=val_set)

model = MyModel()
model = model.cuda()
model = SyncBatchNorn.convert_sync_batchnorm(model)
model = DistributedDataParallel(model, device_ids=[local_rank], output_devices=local_rank, find_unused_parameters=True)

cls_criterion = nn.CrossEntropyLoss()
optim = optimizer()
scheduler = scheduler()

for epoch in range(start_epoch, end_epoch):
    trainer_sampler.set_epoch(epoch)
    train()
    eval()
    if local_rank==0:
        log()

if local_rank==0:
    torch.save(model.module.state_dict())

运行代码：

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py

二、分块总结

2.1 运行代码部分

因为torch.dist对程序的入口有更改，所以先总结运行代码。torch.dist跟DataParallel不同，需要多进程运行程序，以此达到多机训练的效果。在官方的实现中，使用torch.distributed.launch来启动多进程。除此之外，还使用torch.multiprocessing来手动执行多进程程序。在最新的版本中，官方打算将python -m torch.distributed.lanuch用torchun替代，这里暂时不用。

torch.distributed.launch使用方式如下：

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py

其中nproc_per_node表示开启的进程数，这里与使用的卡数保持一致。torch.distributed.launch会自动分配一些参数到主程序中，也可以手动修改。这次关注到的有：RANK表示进程的优先级，也可以认为是进程的序列号；MASTER_ADDR和MASTER_PORT分别表示通讯的地址和端口，torch.distributed.launch会将其设置为环境变量；WORLD_SIZE表示gpu*节点数，本例里就是gpu数量。

2.2 训练代码部分

讲解训练代码的细节，细节还不少。

导入模块

import os
import argparse
import torch
from torch.nn import SyncBatchNorm
from torch.nn.parallel import DistributedDataParallel
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler

argparse：

parser = argparse.ArgumentParser(description='training')
parser.add_argument('--local_rank', type=str, help='local rank for dist')
args = parser.parse_args()

需要承接torch.distributed.launch给main.py传递的local_rank参数。local_rank参数表示进程的优先级，也可以认为是进程的序列号。

环境初始化：

print(os.environ['MASTER_ADDR'])
print(os.environ['MASTER_PORT'])
world_size = torch.cuda.device_count()
local_rank = args.local_rank

dist.init_process_group(backend='nccl')

torch.cuda.set_device(local_rank)

• local_rank表示进程的优先级，也可以认为是进程的序列号；MASTER_ADDR和MASTER_PORT分别表示通讯的地址和端口，torch.distributed.launch会将其设置为环境变量；world_size表示gpu*节点数，本例里就是gpu数量。这里代码里print出来展示。
• dist.init_process_group(backend=‘nccl’)初始化torch.dist的环境。这里backend选择nccl来进行通讯，可以用dist.is_nccl_avaliable()来查看是否可用nccl。除此之外也可以在这里设置一些其他的环境参数。
• torch.cuda.set_device(local_rank)设置环境CUDA序号

数据集设置：

train_dataset = Dataset(...)
train_sampler = DistributedSampler(train_dataset)
train_loader = Dataloader(dataset=train_dataet, sampler=train_sampler, shuffle=False)

val_set = Dataset()
val_loader = Dataloader(dataset=val_set)

• 对训练数据集做修改。将dataloader的sampler修改为DistributedSampler，这样保证其每个进程采样的数据是不同的
• 训练集的dataloader的shuffle只能设置为False，DistributedSampler会进行shuffle，如果dataloader再shuffle的话会打乱次序，导致多进程分配的数据不对
• batch_size设置的是每个进程的，因此不需要像dataparalle一样乘以卡数
• 对验证集可以不做修改，如果每个进程不同的话需要再整合所有进程的结果

模型加载：

model = MyModel()
model = model.cuda()
model = SyncBatchNorn.convert_sync_batchnorm(model)
model = DistributedDataParallel(model, device_ids=[local_rank], output_devices=local_rank, find_unused_parameters=True)

cls_criterion = nn.CrossEntropyLoss()
optim = optimizer()
scheduler = scheduler()

• 跟DataParallel类似的是，加载的模型需要用DDP（DistributedDataParallel）重载一下。这里如果运行时报错有unused_parameters，那么就设置find_unused_parameters=True
• DDP从原理上应该是多机通讯更新梯度从而保证模型的参数都是一样的，而DataParallel则是在一张卡上集中更新模型权重，再复制到其他卡上
• 对于损失函数、优化器和sheduler都不需要DDP，如果有需要更新的参数的话还是需要DDP重载

训练代码：

for epoch in range(start_epoch, end_epoch):
    trainer_sampler.set_epoch(epoch)
    train()
    eval()
    if local_rank==0:
        log()

• sampler.set_epoch(epoch)，这样设置一样才能保证数据是乱序的
• local_rank==0保证只有0进程进行log
• 每个进程都进行eval，我现在不这么做的话会因为其他进程一直没有收到通讯而超时

保存模型：

if local_rank==0:
    torch.save(model.module.state_dict())

• 只需要在0号进程进行保存就行了
• 对DDP用.module

三、运行效果

从我跑的结果来看，训练过程挺好，加速明显（没有展示结果和量化指标）

TODO：

1、可以把验证过程并行化一下加速

问题在于如何合并多进程验证结果。现在想到两套方案：一是用dist的多进程通讯，需要学习dist的其他api；二是将多进程的结果保存到磁盘，再进行合并

2、可以尝试多机多卡