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

　　

　　Nn。DataParallel更容易使用，但是只需要在一台机器上使用。Nn。DataParalllel在每个批处理中，只使用一个进程来计算模型权重并将其分配给每个GPU。 　　

　　

　　在这篇文章中，我将介绍如何nn。数据并行和神经网络。详细的分布式数据并行工作。还会介绍两者的主要区别，以及如何在多个GPU中训练。我们先来看看神经网络训练的工作原理。 　　

　　

　　首先，我们来回顾一下训练神经网络通常是如何工作的。 　　

　　

　　 　　

　　

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

　　

　　正向传递，其中输入由神经网络处理，通过比较预测标签和实际标签来计算损失函数。反向传递完成，根据损耗计算各参数的梯度(使用反向传播)。对于大于1的批量大小，我们可能希望批量规范化培训。 　　

　　

　　DataParallelDataParallel有助于将训练分布到一台计算机上的多个GPU。下面详细介绍一下DataParallel的工作原理。当通过数据并行来训练神经网络时，将发生以下步骤： 　　

　　

　　 　　

　　

　　Mini-batch在GPU: 0上拆分拆分mini-batch，并将其移动到所有不同的GPU上复制机器学习模型到GPUForward pass发生在所有不同的GPU上，根据GPU: 0上的网络输出计算损耗，并将损耗返回给不同的GPU。在每个GPU上计算GPU: 0上的梯度，并使用优化器更新GPU: 0上的模型一个简单的例子。 　　

　　

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

　　

　　 　　

　　

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

　　

　　 　　

　　

　　第4-14行：定义这个神经网络的层。 　　

　　

　　第16-21行：定义向前传球 　　

　　

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

　　

　　 　　

　　

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

　　

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

　　

　　第2445行：这是用于训练神经网络的循环。我们加载输入和预期输出。我们运行向前传递和向后传递以及优化器。 　　

　　

　　“nn”的分布式数据并行。DistributedDataParallel”，这台计算机的每个GPU都有一个进程，每个机器学习模型都由每个进程控制。GPU可以都在同一个节点上，也可以跨多个节点。进程/GPU之间只传递渐变。 　　

　　

　　 　　

　　

　　在训练过程中，每个进程从磁盘加载自己的小批量，并将其传递给GPU。每个GPU都会向前传递，然后所有GPU上的梯度都会减小。每一层的梯度都不依赖于前一层，所以梯度all-reduce与backwards pass同时计算，进一步缓解网络瓶颈。在反向传递结束时，每个节点都有一个平均梯度，从而确保机器学习模型的权重同步。 　　

　　

　　教程 　　

　　

　　为此，我们需要一个Python脚本来为每个GPU启动一个进程。每个进程都需要知道使用哪个GP。 　　

U，以及它在所有正在运行的进程中的位置。我们需要在每个节点上运行脚本。

　　

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

　　

　　

我们来看一下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计算损失。

　　

　　

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

　　

减小batch_size使用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。