pytorch指定用多张显卡训练_如何在多GPU上训练PyTorch模型

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

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

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

训练

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

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

forward pass，其中输入由神经网络处理通过将预测标签与真实标签进行比较来计算损失函数完成backward pass，根据损失计算每个参数的梯度(使用反向传播)使用梯度更新参数对于大于1的batch sizes，我们可能希望对训练进行批归一化。

DataParallel

DataParallel帮助将训练分散到一台计算机上的多个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行：定义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计算损失。

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

减小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。