最近需要训练OCR通用识别模型，单节点训练收敛一次大概需要一周。为了提升训练的效率，我就开始尝试在公司的平台上使用分布式的方式去训练，确实训练速度有很大的提升，这边主要介绍如何修改单节点训练代码，快速地让代码在分布式环境运行起来。如果想看官方的示例可以点 。


官方示例pytorch.org

环境配置

具体的包可以根据个人的需求去安装，但是有一点要注意的事，需要保证每个机器安装的cuda版本一致，python版本一致，pytorch版本一致。之前就遇到过一个坑，某个节点的cuda版本过低，导致如果分配到该节点，代码执行就会失败。

还有比较重要的一点就是需要给NCCL设置环境变量，你可以用ifconfig去查询当前机器私有IP的网络接口名称，我们公司平台上所有机器都是eno1，所有我只要这样设置就可以了。



export NCCL_SOCKET_IFNAME=eno1



训练代码修改

初始化

pytorch分布式训练是以单张显卡为一个进程相互通信，比如你的资源配置是两个节点，每个节点四张显卡，那你就是需要起8个进程（假设单卡可以满足当前的模型），所以每个进程都需要知道这四件事情：

1.一共有多少个进程一起干活。

2.他自己编号是多少。

3.他使用哪张显卡。

4.统一接收分发消息的进程的地址。

这些在初始化的时候就需要指定。



dist.init_process_group(backend=dist_backend, init_method=dist_url, rank=rank, world_size=world_size)
torch.cuda.set_device(local_rank)
assert torch.distributed.is_initialized()



其中dist_backend一般设置为'nccl'，dist_url的格式为”tcp://{ip}:{端口号}“（例如"tcp://172.31.22.234:23456"）

上面的代码就明确了四件事情

  1. 一共world_size个进程在干活。
  2. 当前的进程编号是rank。
  3. 当前进程使用的显卡编号为local_rank。
  4. dist_url就是大家通信的地址，要注意的是rank为0的进程一定在这个地址对应的机器上运行。

模型代码修改



model = DistributedDataParallel(model, device_ids=[opt.local_rank], output_device=opt.local_rank)



使用torch.nn.parallel中的DistributedDataParallel模块来负责分布式训练中的数据处理。其中入参model是继承自nn.Module类型的实例。

Dataloader代码修改

为了支持分布式的训练，需要把torch.utils.data.distributed模块中的DistributedSampler对象传给torch.utils.data.DataLoader。



train_sampler = torch.utils.data.distributed.DistributedSampler(_dataset)
_data_loader = torch.utils.data.DataLoader(
               _dataset, batch_size=_batch_size,
               shuffle=train_sampler is None,
               num_workers=int(opt.workers),
               sampler=train_sampler)



如果你的代码对DataLoader有做封装，只要在实例化DataLoader的时候使用指定的sampler即可。

关于训练

训练流程代码与单节点训练任务无异，要注意的是，在你存储模型权重，打印训练log，验证当前模型在验证集上的结果时，需要指定某一个rank号的进程去做，不然所有的进程可能会把相同的事情都做一边，比如你会发现你的log文件里有多份相同的logs。