梯度累积（gradient accumulation）

Gradient Accumulation

在内存受限的情况下，始终建议以较小的批大小训练神经网络。

尽管如此，有一种叫做梯度累积（gradient accumulation）的技术，可以让我们在不必显式增大批大小的情况下（逻辑上）增大批大小。感到困惑？让我们在本章中加以理解。但在此之前，我们必须先了解……

神经网络中的主要内存开销

来自于反向传播。这是因为在反向传播期间，我们必须将各层激活值存储在内存中。毕竟，它们被用来计算梯度。

网络越大，需要存储在内存中的激活值就越多。而且，在内存受限的情况下，使用较大的批大小会导致：

• 存储大量激活值
• 使用这些大量激活值来计算梯度

这可能导致资源消耗超过可用量——从而导致训练失败。但通过减小批大小，我们可以限制内存使用并完成网络训练。

假设我们正在用小批量训练一个神经网络。

我们按如下方式训练网络：

• 在每个小批量上：
  • 运行前向传播并存储激活值。
  • 在反向传播期间：■ 计算损失 ■ 计算梯度 ■ 更新权重

梯度累积修改了反向传播的最后一步，即权重更新。更具体地说，我们不必在每个小批量上都更新权重，而是可以这样做：

1. 在一个小批量上运行前向传播。2. 计算网络中权重的梯度值。3. 暂不更新权重。4. 在下一个小批量上运行前向传播。5. 计算权重的梯度值，并将其与第 2 步得到的梯度相加。6. 对更多小批量重复步骤 3-5。7. 仅在处理了几个小批量之后才更新权重。

这一技术之所以有效，是因为跨多个小批量累积

梯度所得到的梯度之和，与我们把它们放在一起处理时的梯度之和相同。因此，从逻辑上讲，使用梯度累积，我们可以在不必显式增大批大小的情况下模拟更大的批大小。

例如，假设我们想使用 64 的批大小。然而，当前内存只能支持 16 的批大小。没问题！

• 我们可以使用大小为 16 的批大小。
• 我们可以累积每个小批量的梯度。
• 我们可以每 8 个小批量才更新一次权重。

因此，实际上我们使用的是 16*8（=128）的批大小，而不是我们最初打算的 64。

让我们来看看如何实现这一点。在 PyTorch 中，典型的训练循环实现如下：

• 我们清零梯度
• 运行前向传播
• 计算损失
• 计算梯度
• 更新权重

然而，如前所述，如有需要，我们可以仅在几次迭代之后才更新权重。因此，我们必须继续累积梯度，而这正是 loss.backward() 所做的事情。

此外，由于 optimizer.zero_grad() 会清零梯度，我们必须仅在更新权重之后才执行它。这一思路的实现如下：

• 首先，我们定义 acc_steps——即希望在其之后更新权重的小批量数量。
• 接下来，我们运行前向传播。
• 然后，我们计算损失和梯度。
• 如前所述，我们暂时不更新权重，而是让梯度再累积几个小批量。
• 我们仅在 if 条件为真时才更新权重。
• 更新之后，我们清除已累积的梯度。

这样，我们就能在内存受限的环境下优化神经网络训练。

在结束之前，必须指出：梯度累积并不是在内存受限情况下改善运行时间的灵丹妙药。事实上，我们也可以从我的实验中验证这一点：

相反，它的目标是降低整体内存使用。当然，确实我们只在几次迭代之后才更新权重，所以它会比每次迭代都更新稍快一些。然而，我们仍然在处理和计算小批量上的梯度，这才是这里的核心操作。尽管如此，好处在于即使你不受内存限制，也仍然可以使用梯度累积。

• 指定你通常使用的批大小。
• 运行前向传播。
• 计算损失和梯度。
• 仅在几次迭代之后才更新。

你可以在这里下载 notebook：https://bit.ly/3xNCfFt。