梯度检查点（gradient checkpointing）

Gradient Checkpointing

神经网络主要通过两种方式使用内存：

1. 当它们存储模型权重时（这是固定的内存占用）。2. 当它们被训练时（这是动态的）。它以两种方式发生：a. 在前向传播期间，计算并存储所有层的激活值（activation）。b. 在反向传播期间，计算每一层的梯度。后者，即动态内存占用，常常限制我们用更大的批大小（batch size）训练更大的模型。

这是因为内存占用与批大小成比例增长。话虽如此，有一种相当出色的技术可以让我们在保持总体内存占用不变的同时增大批大小。它被称为梯度检查点（gradient checkpointing），根据我的经验，这是一种被严重低估的、用于降低神经网络内存开销的技术。让我们更详细地了解一下它。

梯度检查点基于对神经网络典型工作方式的两个关键观察：

1. 某一层的激活值可以仅由上一层的激活值计算得到。例如，在下图中，“B 层”的激活值可以仅由“A 层”的激活值计算得到。

2. 更新一层的权重只依赖两件事：a. 该层的激活值。b. 在下一（右侧）层中计算出的梯度（或者更确切地说是运行中的梯度）。

梯度检查点正是利用这两个观察来优化内存占用。它的工作原理如下：

• 步骤 1) 在前向传播之前，将网络划分为若干段：
• 42

• 步骤 2) 在前向传播期间，只存储每段第一层的激活值。当其余激活值已被用于计算下一层的激活值后，就将其丢弃。
• 步骤 3) 现在进入反向传播（backpropagation）。要更新某一层的权重，我们需要它的激活值。因此，我们使用该段第一层重新计算这些激活值。例如，如下图所示，为了更新红色层的权重，我们利用已经存在于内存中的青色层的激活值来重新计算它们的激活值。
• 完成！这就是梯度检查点的工作方式。

总结一下，其思想是我们不必在内存中存储所有中间激活值。相反，只存储其中一部分，并在需要时才重新计算其余部分，可以显著降低内存需求。整个思路在直觉上也是合理的。

这也让我们能够在更大的数据批次上训练网络。通常情况下，我的观察是梯度检查点（GCP）可以将内存使用量至少降低 50-60%，这是相当可观的。

当然，由于我们要计算两次某些激活值，这确实以运行时间增加为代价，通常增加幅度在 15-25% 之间。因此，内存与运行时间之间始终存在权衡。

话虽如此，它的另一个优势是允许我们使用更大的批大小，这可以稍微（虽然不能完全）抵消增加的运行时间。尽管如此，梯度检查点仍是训练更大模型的一种极其强大的技术，我发现它在某些时候相当有用，而无需诉诸分布式训练等更密集的技术。令人欣慰的是，梯度检查点也已被 PyTorch 等许多开源深度学习框架实现。这里有一个演示。

要使用它，我们首先导入必要的库和函数：

接下来，我们定义我们的神经网络：

如上所示，在 forward 方法中，我们使用 checkpoint_sequential 方法来使用梯度检查点，并将网络划分为两段。接下来，我们就可以像通常那样进行网络训练了。