apesv100数据库_PyTorch 单机多卡操作总结：分布式DataParallel，混合精度，Horovod)

点击上方“机器学习与生成对抗网络”，关注"星标"

获取有趣、好玩的前沿干货！

作者丨科技猛兽@知乎  编辑丨极市平台

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

在上一篇文章中(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 进行一系列复杂的CPU、GPU分配任务，PyTorch为我们提供了一个很方便的启动器

torch.distributed.launch 用于启动文件，所以我们运行训练代码的方式就变成了这样：

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

其中的

--nproc_per_node 参数用于指定为当前主机创建的进程数，由于我们是单机多卡，所以这里node数量为1，所以我们这里设置为所使用的GPU数量即可。

2.2.2 初始化

在启动器为我们启动python脚本后，在执行过程中，启动器会将当前进程的(其实就是 GPU的)index 通过参数传递给 python，我们可以这样获得当前进程的 index：即通过参数

local_rank 来告诉我们

当前进程使用的是哪个GPU，用于我们在每个进程中指定不同的device：

def parse():

parser = argparse.ArgumentParser()

parser.add_argument('--local_rank', type=int, default=0，help='node rank for distributed training')

args = parser.parse_args()

return args

def main():

args = parse()

torch.cuda.set_device(args.local_rank)

torch.distributed.init_process