PyTorch 单机多卡操作总结：分布式DataParallel，混合精度，Horovod)-技术圈

↑ 点击蓝字关注极市平台

作者丨科技猛兽@知乎

来源丨https://zhuanlan.zhihu.com/p/158375055

编辑丨极市平台

极市导读

本文介绍了数种实现单机多卡操作的方法，含有大量代码，并给出了实践中作者踩过的坑及其解决方案。>>加入极市CV技术交流群，走在计算机视觉的最前沿

在上一篇文章中（https://zhuanlan.zhihu.com/p/158375254）我们看到了多GPU训练，也就是最简单的单机多卡操作nn.DataParallel。但是很遗憾这种操作还不够优秀，于是就有了今天这篇文章~

写这篇文章的时候看了很多的tutorials，附在文末了，在此先向文末的每位作者致敬，感谢大佬们！

其实单机多卡的办法还有很多(如下)，而且上篇的方法是相对较慢的。

1、nn.DataParallel 简单方便的 nn.DataParallel

2、torch.distributed 使用 torch.distributed 加速并行训练

3、apex 使用 apex 再加速。

这里，记录了使用 4 块 Tesla V100-PICE 在 ImageNet 进行了运行时间的测试，测试结果发现 Apex 的加速效果最好，但与 Horovod/Distributed 差别不大，平时可以直接使用内置的 Distributed。Dataparallel 较慢，不推荐使用。那好像上一篇白写了~

看到这里你可能已经懵逼了，莫慌，下面会分别进行介绍(这里先附上一篇教程，看不懂的话直接当做没看见即可)：

https://yangkky.github.io/2019/07/08/distributed-pytorch-tutorial.html

1. 先问两个问题

问1：为啥非要单机多卡？

答1：加速神经网络训练最简单的办法就是上GPU，如果一块GPU还是不够，就多上几块。

事实上，比如BERT和GPT-2这样的大型语言模型甚至是在上百块GPU上训练的。

为了实现多GPU训练，我们必须想一个办法在多个GPU上分发数据和模型，并且协调训练过程。

问2：上一篇讲得单机多卡操作nn.DataParallel，哪里不好？

答2：要回答这个问题我们得先简单回顾一下nn.DataParallel，要使用这玩意，我们将模型和数据加载到多个 GPU 中，控制数据在 GPU 之间的流动，协同不同 GPU 上的模型进行并行训练。具体怎么操作？

我们只需要用 DataParallel 包装模型，再设置一些参数即可。需要定义的参数包括：

参与训练的 GPU 有哪些，device_ids=gpus。
用于汇总梯度的 GPU 是哪个，output_device=gpus[0] 。

DataParallel 会自动帮我们将数据切分 load 到相应 GPU，将模型复制到相应 GPU，进行正向传播计算梯度并汇总：

model = nn.DataParallel(model.cuda(), device_ids=gpus, output_device=gpus[0])

值得注意的是，模型和数据都需要先 load 进 GPU 中，DataParallel 的 module 才能对其进行处理，否则会报错：

# main.py
import torch
import torch.distributed as dist

gpus = [0, 1, 2, 3]
torch.cuda.set_device('cuda:{}'.format(gpus[0]))

train_dataset = ...

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=...)

model = ...
model = nn.DataParallel(model.to(device), device_ids=gpus, output_device=gpus[0])

optimizer = optim.SGD(model.parameters())

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
      images = images.cuda(non_blocking=True)
      target = target.cuda(non_blocking=True)
      ...
      output = model(images)
      loss = criterion(output, target)
      ...
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

稍微解释几句：model.to(device)将模型迁移到GPU里面，images.cuda，target.cuda把数据迁移到GPU里面。
nn.DataParallel(model.to(device), device_ids=gpus, output_device=gpus[0])包装模型。

缺点：

在每个训练批次（batch）中，因为模型的权重都是在一个进程上先算出来，然后再把他们分发到每个GPU上，所以网络通信就成为了一个瓶颈，而GPU使用率也通常很低。
除此之外，nn.DataParallel 需要所有的GPU都在一个节点（一台机器）上，且并不支持 Apex 的混合精度训练。

一句话，一个进程算权重使通信成为瓶颈，nn.DataParallel慢而且不支持混合精度训练。

2. 使用 torch.distributed 加速并行训练：

DataParallel：单进程控制多 GPU。

DistributedDataParallel：多进程控制多 GPU，一起训练模型。

2.1 介绍

在 1.0 之后，官方终于对分布式的常用方法进行了封装，支持 all-reduce，broadcast，send 和 receive 等等。通过 MPI 实现 CPU 通信，通过 NCCL 实现 GPU 通信。官方也曾经提到用 DistributedDataParallel 解决 DataParallel 速度慢，GPU 负载不均衡的问题，目前已经很成熟了。

与 DataParallel 的单进程控制多 GPU 不同，在 distributed 的帮助下，我们只需要编写一份代码，torch 就会自动将其分配给n个进程，分别在n个 GPU 上运行。

和单进程训练不同的是，多进程训练需要注意以下事项：

在喂数据的时候，一个batch被分到了好几个进程，每个进程在取数据的时候要确保拿到的是不同的数据（DistributedSampler）；
要告诉每个进程自己是谁，使用哪块GPU（args.local_rank）；
在做BatchNormalization的时候要注意同步数据。

2.2 使用方式

2.2.1 启动方式的改变

在多进程的启动方面，我们不用自己手写 multiprocess 进行一系列复杂的C