GPU如何训练大批量模型？方法在这里

深度学习模型和数据集的规模增长速度已经让 GPU 算力也开始捉襟见肘，如果你的 GPU 连一个样本都容不下，你要如何训练大批量模型？通过本文介绍的方法，我们可以在训练批量甚至单个训练样本大于 GPU 内存时，在单个或多个 GPU 服务器上训练模型。

分布式计算

2018 年的大部分时间我都在试图训练神经网络时克服 GPU 极限。无论是在含有 1.5 亿个参数的语言模型（如 OpenAI 的大型生成预训练 Transformer 或最近类似的 BERT 模型）还是馈入 3000 万个元素输入的元学习神经网络（如我们在一篇 ICLR 论文《Meta-Learning a Dynamical Language Model》中提到的模型），我都只能在 GPU 上处理很少的训练样本。

但在多数情况下，随机梯度下降算法需要很大批量才能得出不错的结果。

如果你的 GPU 只能处理很少的样本，你要如何训练大批量模型？

有几个工具、技巧可以帮助你解决上述问题。在本文中，我将自己用过、学过的东西整理出来供大家参考。

在这篇文章中，我将主要讨论 PyTorch 框架。有部分工具尚未包括在 PyTorch（1.0 版本）中，因此我也写了自定义代码。

我们将着重探讨以下问题：

在训练批量甚至单个训练样本大于 GPU 内存，要如何在单个或多个 GPU 服务器上训练模型；

如何尽可能高效地利用多 GPU 机器；

在分布式设备上使用多个机器的最简单训练方法。

在一个或多个 GPU 上训练大批量模型

你建的模型不错，在这个简洁的任务中可能成为新的 SOTA，但每次尝试在一个批量处理更多样本时，你都会得到一个 CUDA RuntimeError：内存不足。

这位网友指出了你的问题！

但你很确定将批量加倍可以优化结果。

你要怎么做呢？

这个问题有一个简单的解决方法：梯度累积。

梯度下降优化算法的五个步骤。

与之对等的 PyTorch 代码也可以写成以下五行：

predictions = model(inputs)               # Forward pass
loss = loss_function(predictions, labels) # Compute loss function
loss.backward()                           # Backward pass
optimizer.step()                          # Optimizer step
predictions = model(inputs)               # Forward pass with new parameters

在 loss.backward() 运算期间，为每个参数计算梯度，并将其存储在与每个参数相关联的张量——parameter.grad 中。

累积梯度意味着，在调用 optimizer.step() 实施一步梯度下降之前，我们会对 parameter.grad 张量中的几个反向运算的梯度求和。在 PyTorch 中这一点很容易实现，因为梯度张量在不调用 model.zero_grad() 或 optimizer.zero_grad() 的情况下不会重置。如果损失在训练样本上要取平均，我们还需要除以累积步骤的数量。

以下是使用梯度累积训练模型的要点。在这个例子中，我们可以用一个大于 GPU 最大容量的 accumulation_steps 批量进行训练：

model.zero_grad()                                   # Reset gradients tensors
for i, (inputs, labels) in enumerate(training_set):
    predictions = model(inputs)                     # Forward pass
    loss = loss_function(predictions, labels)       # Compute loss function
    loss = loss / accumulation_steps                # Normalize our loss (if averaged)
    loss.backward()                                 # Backward pass
    if (i+1) % accumulation_steps == 0:             # Wait for several backward steps
        optimizer.step()                            # Now we can do an optimizer step
        model.zero_grad()                           # Reset gradients tensors
        if (i+1) % evaluation_steps == 0:           # Evaluate the model when we...
            evaluate_model()                        # ...have no gradients accumulated

扩展到极致

你可以在 GPU 上训练连一个样本都无法加载的模型吗？

如果你的架构没有太多跳过连接，这就是可能的！解决方案是使用梯度检查点（gradient-checkpointing）来节省计算资源。