Training Deep Nets with Sublinear Memory Cost 笔记

摘要

我们提出了一种减少深度神经网络训练时内存消耗的系统性方法。具体来说，我们设计了一个算法，训练一个 $n$ 层网络仅耗费 $O(\sqrt{n})$ 的内存，每个mini-batch只需要一个额外的前向计算成本。由于许多最先进的模型已经达到了GPU显存的上限，我们的算法允许探索更深入更复杂的模型，并有助于推进深度学习研究的创新。我们专注于降低训练期间存储中间特征图和梯度的内存成本。计算图分析用于自动原地操作和内存共享优化。我们证明通过一种内存利用更高效的训练算法以及一点额外的计算成本，可以通过计算来换取内存资源。在极端情况下，我们的分析还表明，内存消耗可以减少到 $O(\log n)$ ，而正向计算的额外开销只需 $O(n\log n)$ 。我们的实验表明，在ImageNet上我们可以将1000层深度残差网络的内存成本从48G降低到7G。类似地，在非常长的序列中训练复杂递归神经网络时亦可显着的降低内存成本。

介绍

在本文中，我们提出了一种减少深度神经网络训练时内存消耗的系统性方法。我们主要关注于减少存储中间结果（特征映射）和梯度的内存成本，因为在许多常见深度架构中，与中间特征映射的大小相比，参数的大小相对较小。我们使用计算图分析来执行自动原地操作和内存共享优化。更重要的是，我们提出了一种新的以计算交换内存的方法。从而，我们给出了一个实用的算法，其花费 $O(\log n)$ 的内存用于特征映射，仅以双倍正向计算成本训练n层网络。有趣的是，我们还表明，在极端情况下，可以使用 $O(\log n)$ 的内存存储特征映射来训练一个 $n$ 层网络。

我们最近见证了深度神经网络在许多领域的成功[8]，如计算机视觉、语音识别、自然语言处理和强化学习。深度神经网络新架构方面的创新带来了许多成功。卷积神经网络[15,14,13,10]对空间模式进行建模并给出计算机视觉任务中的顶尖结果。递归神经网络，如长期短期记忆[12]，在序列建模和结构预测中显示出鼓舞人心的结果。这些新模型的一个共同趋势是使用更深的架构[18,14,13,10]来捕捉大量训练数据中的复杂模式。由于存储特征映射及其梯度的成本随网络深度线性变化，因此我们探索深层模型的能力受到设备（通常为GPU）内存的限制。例如，如[11]所述，当前最先进的一个模型已经耗尽了内存。从长远来看，理想的机器学习系统应该能够从越来越多的训练数据中不断学习。由于最佳模型大小和复杂性通常随着训练数据增多而增加，因此节约内存的训练算法是非常重要的。

减少内存消耗不仅使我们能够训练更大的模型。它还可以实现更大的批量大小，从而提高批量操作符的设备利用率和稳定性，如批量归一化[13]。对于内存有限的设备，它有助于改善内存局部性，并可能产生更好的内存访问模式。它还使我们训练深度卷积神经网络时能够从模型并行性切换到数据并行性，这在某些情况下可能是有益的。我们的解决方案使我们能够训练更深的卷积神经网络，以及具有更大展开步长的递归神经网络。结合本文提出的内存优化技术，我们为深度学习框架提供指导。我们还将公开我们内存优化算法的实现。

基于计算图的内存优化

我们首先回顾一下计算图的概念和内存优化技术。其中一些技术已经被现有的框架如Theano[5,4]，Tensorflow[1]和MXNet[6]所使用。计算图由表示操作之间依赖关系的操作节点和边组成。图1给出了一个双层全连接神经网络计算图的例子。这里我们使用粗粒度的向前和向后操作来使图更简单。我们通过隐藏权重节点和权重的梯度来进一步简化图。在实践中使用的计算图可能更复杂，并且混合包含粗/细粒度操作。本文提出的分析可以直接用于那些更一般的情况。

图1 两层全连接神经网络训练过程的计算图和可能的内存分配计划。每个节点代表一个操作，每个边表示操作之间的依赖关系。具有相同颜色的节点共享内存以存储输出或反向传播的梯度。为使图更清楚，我们省略了图中权重及其输出梯度节点，并假定在后向运算时会计算权重的梯度。我们还注释了使用原地和共享策略的两个地方。

一旦给出网络配置（前向图），我们就可以构建相应的后向通道进行梯度计算。可以通过以逆向拓扑顺序遍历配置来构建后向路径，并且像在正常反向传播算法中那样应用后向运算符。图1中的后向路径显式给出了梯度计算步骤，使得训练中的梯度计算步骤简化为整个计算图表（包括梯度计算路径）上的正向通过。显式梯度路径还提供了一些其他好处（例如，能够计算更高阶梯度），这超出了本文所讨论的范围。

在训练深度卷积/循环网络时，通常使用大部分内存来存储中间输出和梯度。这些中间结果中的每一个都对应于图中的一个节点。智能分配算法能够通过尽可能内存共享来为这些节点分配最少量的内存。图1展示了例子中两层神经网络可能的分配计划。可以使用两种类型的内存优化

• 原地操作：直接将输出值存储到输入值的内存中。
• 内存共享：回收不再需要的中间结果使用的内存并用于另一个节点。

图1中的分配计划包含两种情况的例子。第一个sigmoid转换使用原地操作来节省内存，而其后向操作又进行了重用。softmax梯度可以与第一个全连接层的梯度共享内存。特定情况下应用这些优化可能会导致错误。例如，如果一个操作的输入仍然是另一个操作所需要的，那么对输入进行原地操作将导致错误的结果。

我们只能在生命周期不重叠的节点之间共享内存。有多种方法可以解决这个问题。一个选择是以每个变量作为节点，变量之间的重叠生命周期为边，构造冲突图，然后运行图着色算法。这将花费 $O(n^2)$ 的计算时间。我们采用更简单的启发式方法，仅需 $O(n)$ 的时间。该算法如图2所示。它以拓扑顺序遍历图，并使用计数器来指示每条记录的活性。当没有其他操作等待其输入时，可能会产生原地操作。另一个节点使用回收标签时会进行内存共享。这也可以作为图遍历的动态运行时算法，并使用垃圾回收器来回收过时的内存。我们将其用作静态内存分配算法，在执行开始之前为每个节点分配内存，以避免在运行时垃圾收集的开销。

图2 计算图上的内存分配算法。每个节点都关联一个活性计数器，用于记录要填充的操作。时间标记用于指示内存共享。如果当前操作是剩下的唯一操作（输入的计数器等于1），则可以执行原地操作。当节点计数器变为零时，可以回收节点的标签。

深度学习框架指南 正如我们从图2中的算法演示图中所看到的那样。数据依赖性导致每个输出的使用期更长，并增加了大型网络的内存消耗。对于深度学习框架来说非常重要是

• 以最小方式声明梯度运算符的依赖性要求。
• 对依赖性信息应用活性分析并启用内存共享。

声明最小的依赖关系很重要。例如，如果sigmoid-backward也取决于第一个fullc-forward的输出，则图1中的分配计划将不可能实现。依赖性分析通常可以将 $n$ 层网络的深度网络预测的内存占用从 $O(n)$ 减少到接近 $O(1)$ ，因为可以在每个中间结果之间进行共享。该技术也有助于减少训练的内存占用，尽管只是一个常数因素。

以计算换内存

通用方法

上一节介绍的技术可以减少深度神经网络训练和预测的内存占用。然而，由于大多数梯度算子都依赖于正向通路的中间结果，因此训练 $n$ 层卷积网络或序列长度为 $n$ 的递归神经网络，我们仍需要 $O(n)$ 的内存用于中间结果。为了进一步减少内存，我们建议删除一些中间结果，并在需要时借助额外的前向计算恢复它们。

更具体地说，在反向传播阶段，我们可以通过从最近的记录结果运行前向计算来获得丢弃的中间结果。为了更清楚地呈现这个想法，我们在Alg. 1 中展示了一个直链前馈神经网络的简化算法。具体来说，神经网络被分为几个部分。该算法只记住每段的输出并删除段内的所有中间结果。反向传播期间在段内重新计算丢弃的结果。因此，我们的内存开销仅需覆盖所有段的输出以及逐段反向传播时的最大内存占用。

只要我们可以将图分成几部分，Alg. 1 也可以推广到其他常用的计算图。但是，直接应用ALG. 1有两个缺点：

• 用户必须手动划分图并编写定制的训练循环;
• 我们无法从上节中介绍的其他内存优化中受益。

我们通过引入基于相同思想的通用梯度图构造算法来解决这一问题。该算法在Alg. 2中给出。在该算法中，用户在计算图的节点上指定函数 $m:\mathcal{V}\to \mathbb{N}$ ，以指示可以重新计算结果的次数。我们将 $m$ 称之为镜像计数函数，因为重新计算本质上是复制（镜像）节点。当所有镜像计数设置为0时，算法退化为正常梯度图。在Alg. 2中指定重新计算模式时 ，用户只需要设置段内节点的 $m(v)=1$ ，段输出节点的 $m(v)=0$ 。镜像计数也可以大于1，这会导致递归泛化，在后续节中会进行讨论。图3显示了内存优化梯度图的一个例子。重要的是，Alg. 2还输出用于计算的遍历顺序，所以可以优化内存使用。此外，这种遍历顺序可以帮助依赖于运行时分配的框架引入控制流依赖项。

丢弃低运算成本操作的结果

通用方法的一个快速应用是丢弃低成本操作的结果并保持计算耗时的结果。这通常用于卷积神经网络中的Conv-BatchNorm-Activation流水线。我们总是可以保留卷积结果，但是丢弃归一化、激活和池化的结果。在实践中，这将节省内存，并且计算开销很小，因为批量归一化和激活函数的计算成本很低。

$O(\sqrt{n})$ 内存代价算法

Alg. 2提供了一种计算换内存的通用方法。但它仍然需要了解应该保留哪个中间结果，以及要重新计算哪个结果。假设我们将 $n$ 层网络划分为 $k$ 段，训练此网络的内存成本如下给出。
$$\mathrm{cost}-\mathrm{total}=\underset{i=1,\dots ,k}{\max }\mathrm{cost}-\mathrm{of}-\mathrm{segment}(i)+O(k)=O\left(\frac{n}{k}\right)+O(k)$$

等式的第一部分是在每个段上反向传播的内存开销。鉴于段是平等分配的，这将转化为 $O(n/k)$ 的成本。等式的第二部分是存储分段之间的中间输出的成本。设置 $k=\sqrt{n}$ ，我们得到 $O(2\sqrt{n})$ 的成本。该算法只需在训练过程中增加一次前向，但将内存开销降低为次线性。由于后向操作几乎是前向的两倍，所以只会稍许减慢计算速度。

图3 内存优化梯度图生成示例。镜像正向路径以表示梯度计算时发生的重计算。用户指定镜像因子以控制是否应放弃或保留结果。

一般在多数情况下，每层的内存成本是不一样的，所以我们不能简单地设置 $k=\sqrt{n}$ 。但是，中间输出和每个阶段成本之间的权衡仍然存在。在这种情况下，给定每个段内的内存开销预算作为单一参数$ B $ ，我们使用Alg. 3以贪婪算法进行分配。不同的$ B $ 给出我们不同的分配计划，要么给中间输出分配更多的内存，要么给每个阶段分配更多的计算。当我们进行静态内存分配时，给定每个分配计划，我们可以得到精确的内存消耗。我们可以使用这些信息对 $ B $ 进行启发式搜索，以找到平衡两者成本的最佳内存方案。搜索步骤的细节在补充材料中介绍。我们发现这种方法在实践中运作良好。我们还可以推广该算法，通过考虑执行每个操作的成本以尽可能地保留耗时的操作。

更一般化的视角：递归和子程序

在本节中，我们提供了上述内存优化方案的另一种理解。具体来说，我们可以将每个段视为一个块操作符，它将段中的所有操作组合在一起。这一思路如图4所示。组合运算符在描述其内部计算的子图上运行来计算梯度。这一观点允许我们将一系列操作视为子程序。子图中的优化不会影响外部世界。因此，我们可以递归地将我们的内存优化方案应用于每个子图。

图4 内存分配优化的递归视图。分段可以看作是一个单一的运算符，它将段内的所有运算符组合在一起。在每个运算符中，执行一个子图来计算梯度。

用递归进一步减少内存 设 $g(n)$ 为在 $n$ 层神经网络上进行前向和后向传播的内存代价。假设我们将 $k$ 个中间结果存储在图中，并在子路径上进行前向和后向传播时递归地应用相同的策略。我们有以下递归公式。
$$g(n)=k+g\left(n/(k+1)\right)$$
求解这个递归公式给了我们
$$g(n)=k{\log }_{k+1}(n)$$
作为一个特例，如果我们设 $k=1$，我们得到$g(n)={\log }_{2}n$。这是一个有趣的结论，因为训练一个 $n$ 层神经网络，现有的所有实现都仅需 $O(n)$ 的内存用于特征映射。这将需要 $O({\log }_{2}n)$ 的正向成本，因而不太可能普遍使用。但它演示了我们如何通过递归换取内存。

深度学习框架指南

在本节中，我们已经证明能够以计算换内存并将其与上一节中提出的系统优化结合起来。这有助于深度学习框架

• 启用选项以删除低成本操作的结果。
• 提供规划算法以提供有效的内存方案。
• 使用户能够在计算图中设置镜像属性以进行内存优化。

虽然最后一个选项不是必需的，但提供这样的接口使得用户可以创造他们自己的内存优化器，并鼓励将来对相关方向进行研究。在这一愿景下，我们支持定制图镜像方案，并将公开源代码。

B预算搜索

Alg. 3使我们可以在给定单个参数 $ B $ 的情况下生成优化的内存方案。该算法依靠近似内存估计获得更快的速度。在方案完成后，我们可以使用静态分配算法计算确切的内存成本。然后，我们可以在 $ B $上进行网格搜索以找到一个好的内存方案。

为了获得网格的设置，我们首先以 $ B = 0 $ 运行分配算法，然后改用 $B=\sqrt{xy}$ 运行一次。这里， $ x $ 和 $ y $ 是 Alg. 3 在第一次运行中的输出。 $ x $ 是存储段间特征映射的近似成本，$ y $ 是运行每段的内存近似成本。 $B=\sqrt{xy}$ 为每一阶段内存成本的估计。至此已经可以提供一个很好的内存方案。然后，我们围绕 $B=\sqrt{xy}$ 设置网格，以进一步优化解决方案。

在实践中，我们发现在 $[B/\sqrt{2},\sqrt{2}B]$ 上使用大小为 $ 6 $ 的网格可以给出实验中良好的内存方案。我们用python实现了分配算法，而未尝试优化速度。我们的代码仅需几秒钟就能获得实验所需的方案。

实验

实验结果及完整内容请参考原文

算法实现

下面是福利时间。对于MXNet用户，作者已经公布代码；而TFBOYS可以使用OpenAI的gradient-checkpointing。

影响力

论文的价值可以通过引用量间接体现。例如引证文献：
MegDet: A Large Mini-Batch Object Detector
Memory-Efficient Implementation of DenseNets