深度学习模型的最大问题之一是它们经常很大，往往无法在单个GPU中进行训练。PyTorch构建了两种在多个GPU中实施分布式训练的方法：nn.DataParalllel和nn.DistributedParalllel。它们是包装和更改代码以及增加在多个GPU中训练网络的功能的简单方法。

nn.DataParallel更容易使用，但只需要在一台机器上使用。nn.DataParalllel在每个批次中，仅使用一个过程来计算模型权重并将其分配给每个GPU。

在这篇文章中，我会详细介绍nn.DataParallel和nn.DistributedDataParalllel是如何工作的。还将介绍两者之间的主要区别以及如何在多个GPU中进行训练。首先来看神经网络训练的工作原理。

训练

首先，让我们复习一下训练神经网络通常是如何工作的。

训练神经网络时，每个循环有四个主要步骤：

1. forward pass，其中输入由神经网络处理
2. 通过将预测标签与真实标签进行比较来计算损失函数
3. 完成backward pass，根据损失计算每个参数的梯度（使用反向传播）
4. 使用梯度更新参数

对于大于1的batch sizes，我们可能希望对训练进行批归一化。

DataParallel

DataParallel帮助将训练分散到一台计算机上的多个GPU中。让我们详细介绍一下DataParallel的工作原理。当使用数据并行训练神经网络时，会发生以下几个步骤:

1. mini-batch在GPU：0上拆分
2. 拆分mini-batch并将其移至所有不同的GPU
3. 将机器学习模型复制到GPU
4. Forward pass发生在所有不同的GPU中
5. 根据GPU：0上的网络输出计算损失，并将损失返回给不同的GPU。在每个GPU上计算梯度
6. 汇总GPU：0上的梯度，并使用优化器更新GPU：0上的模型

一个简单的例子

让我们对此进行Python编码。首先，让我们导入所需的库

我们定义一个非常简单的卷积模型来预测MNIST

第4-14行：定义此神经网络中的层。

第16-21行：定义forward pass

main()函数将接受一些参数并运行训练函数:

第2-6行：我们实例化模型并将其设置为在指定的GPU中运行，并通过使用在多个GPU中并行运行我们的操作DataParallel。

第9-23行：定义损失函数和优化器（我们使用的是SGD）。定义训练数据集（MNIST）和数据的加载器。

第24–45行：这是用于训练神经网络的循环。我们加载输入和预期输出。我们运行forward pass和backward pass以及优化器。

DistributedDataParallel

对于“ nn.DistributedDataParallel”，该计算机每个GPU具有一个进程，并且每个机器学习模型都由每个进程控制。GPU都可以在同一节点上，也可以跨多个节点。在进程/ GPU之间仅传递梯度。

在训练期间，每个进程都从磁盘加载其自己的mini-batch并将其传递到其GPU。每个GPU都会进行forward pass，然后所有GPU上的梯度都会降低。每层的梯度不依赖于先前的层，因此梯度all-reduce与backwards pass同时计算，以进一步缓解网络瓶颈。在backwards pass结束时，每个节点都具有平均梯度，从而确保机器学习模型权重保持同步。

教程

为此，我们需要一个Python脚本来为每个GPU启动一个进程。每个进程都需要知道要使用哪个GPU，以及它在所有正在运行的进程中的位置。我们需要在每个节点上运行脚本。

让我们看看每个函数的变化。

我们来看一下main函数的参数：

• args.nodes 是我们正在使用的节点总数（计算机数量）。
• args.gpus 是每个节点（每台计算机上）上GPU的数量。
• args.nr是当前节点（计算机）在所有节点（计算机）中的rank，从0到args.nodes-1。

让我们逐行浏览新的更改：

第12行：根据节点数和每个节点的GPU，我们可以计算world_size或要运行的进程总数，它等于GPU总数乘以节点数。

第13行：告诉multiprocessing模块要为进程0查找哪个IP地址。它需要此地址，以便所有的进程都能在开始时同步。这需要在所有节点上都是相同的。

第14行：同样，这是查找进程0时要使用的端口。

第15行：我们将产生args.gpus个过程，其中的每一个运行train(i, args)，其中i从0至args.gpus- 1。请记住，请记住，我们在每个节点上运行main()函数，从而总共会有args.nodes* args.gpus= args.world_size个processes。

接下来，让我们看一下对train的修改。

第3行：这是所有进程中进程的global rank。我们将其用于第6行。

第4-6行：初始化进程并与其他进程结合。这是“blocking”，这意味着在所有进程都加入之前，没有任何进程会继续。这里使用NCCL，因为它是最快的。nit_method告诉进程组在何处查找某些设置。在本例中，它查看我们在main中设置的MASTER_ADDR和MASTER_PORT的环境变量。这就是我们将其设置为env://的原因。

第23行：将模型包装为DistributedDataParallel模型。这会将模型复制到每个GPU上。

第35-39行：nn.utils.data.DistributedSampler确保每次加载数据时每个过程都获得不同的训练数据片段。如果要调试并验证每个GPU是否正在加载正确的数据，则可以计算加载到每个GPU中的张量的SHA。

第46和51行：使用nn.utils.data.DistributedSampler而不是shuffling。这就是为什么我们将shuffle设置为false。

例如，要在具有8个GPU的4个节点上运行此程序，我们需要4个终端（每个节点一个）。在节点0上（由main中的第13行设置）：

python src/mnist-distributed.py -n 4 -g 8 -nr 0

然后，在其他节点上：

python src/mnist-distributed.py -n 4 -g 8 -nr i

因为i∈1,2,3。请注意，有效的batch_size现在是per / GPU batch_size（脚本中的值）* GPU的总数。

问题

每当在几个GPU而不是一个GPU中运行相同模型时，可能会出现一些问题。可能发生的最大问题是主GPU可能内存不足。这样做的原因是因为第一个GPU将保存所有不同的输出以供不同的GPU计算损失。

为了解决此问题并减少内存使用量，我们使用两种技术：

1. 减小batch_size
2. 使用Apex实现混合精度

第一种技术非常简单，通常只需更改一个超参数即可。

第二种技术意味着我们将降低神经网络中使用的权重的精度，因此将使用较少的内存。

Apex混合精度

为了解决内存不足的问题，我们建议使用较低精度的数字。这使得我们可以使用更大的批大小，并利用NVIDIA张量核进行更快的计算。

为了使APEX工作，我们需要更改代码的两部分。第一个是在train代码库中的循环内部：

训练步骤

第18行：amp.initialize包装模型和优化器以进行混合精度训练。请注意，在调用之前，模型必须已经在正确的GPU上amp.initialize。该opt_level云从O0，它采用全浮动，通过O3，它使用整个半精度。O1和O2是不同程度的混合精度，有关详细信息，请参见Apex 文档。

第18行:amp.initialize包装了用于混合精度训练的模型和优化器。请注意，在调用amp.initialize之前，模型必须已经在GPU上正确运行。opt_level从O0到使用半精度的O3。O1和O2是不同程度的混合精度，其细节可以在Apex文件中找到。

第20行：apex.parallel.DistributedDataParallel是nn.DistributedDataParallel的替代产品。我们不再需要指定GPU，因为Apex每个进程只允许一个GPU。它还假设脚本torch.cuda.set_device(local_rank)在将模型移至GPU之前调用（第10行）。

37-38行：混合精度训练要求对损失进行缩放，以防止梯度下溢。Apex会自动执行此操作。

请确保无论何时初始化AMP，都要设置opt_level=O1，因为它的实现有一个bug

检查点

每当使用Apex时，我们需要改变保存检查点的方式，并将它们加载到我们的模型中:

第5行：将添加amp.state_dict到检查点

第19行：在这里将state_dict加载到amp。

结论

通过本文，您应该能够开始在多个GPU中训练机器学习模型。我们建议先尝试在一个GPU中训练一个小型模型，然后再尝试将训练扩展到多个GPU。