人工智能研学社

问题：GPU 内存限制

GPU 在深度神经网络训练之中的强大表现无需我赘言。通过现在流行的深度学习框架将计算分配给 GPU 来执行，要比自己从头开始便捷很多。然而，有一件事你会避之唯恐不及，即 GPU 的动态随机存取内存（DRAM（Dynamic Random Access Memory））限制。

在给定模型和批量大小的情况下，atv，事实上你可以计算出训练所需的 GPU 内存而无需实际运行它。例如，使用 128 的批量训练 AlexNet 需要 1.1GB 的全局内存，而这仅是 5 个卷积层加上 2 个全连接层。如果我们观察一个更大的模型，比如 VGG-16，使用 128 的批量将需要大约 14GB 全局内存。目前最先进的 NVIDIA Titan X 内存为 12GB。VGG-16 只有 16 个卷积层和 3 个全连接层，这比大约有 100 个层的 resnet 模型小很多。

　　图 1: 当使用基线、全网分配策略时 GPU 内存的使用情况（左轴）。(Minsoo Rhu et al. 2016)

现在，如果你想要训练一个大于 VGG-16 的模型，而你又没有一个具有内存的 GPU，你也许有几个解决内存限制问题的选择。

减小你的批量大小，但这可能会妨碍你的训练速度和精确度。

在多 GPU 环境下做模型并行，这是另一个复杂的事情。

缩小你的模型，如果你不情愿做出上述两个选择，或者已经尝试但效果不好。

或者你也可以等待下一代更强大的 GPU 出现。然而将来的网络变得更大更深是大势所趋，我们并不希望物理内存的限制成为算法开发的一个障碍。

观察：什么占用内存？

我们可以根据功能性把 GPU 内存中的数据分为 4 个部分：

模型参数

特征图

梯度图

工作区

前 3 个功能容易理解。模型参数的意义对于所有人来说都很熟悉了。特征图是正向过程中生成的中间结果。梯度图是反向过程中生成的中间结果。而工作区是 cuDNN 库函数所使用的临时变量／矩阵的一个缓冲区。对于一些 cuDNN 函数，用户需要将缓冲区作为函数参数传给 GPU 内核函数。一旦函数返回，该缓冲区将被释放。

我们可以看到，一般而言，我们有的层越多，分配给特征图的内存比重就越多（由上图中的三角形表示）。我们也可以看到对于像 VGG-16 这样的更大模型来说，这一比例基本上要超过 50%。

想法：使用 CPU 内存作为临时容器

有一个关于特征图的事实：它们在正向过程中生成之后立刻被用于下一层，并在稍后的反向过程中仅再用一次。每个 GPU 内核函数仅使用与当前层（通常只有 1 个张量）相关的特征映射。这将导致绝大多数内存在几乎所有的时间上出现空置的情况（它们保有数据但不使用）。

这一想法是：如果 GPU 内存中的大部分数据出现空置，为什么不把它们保存在更便宜的 CPU 内存上呢？下图更清晰地展现了这一想法。

左侧部分所示的间隙表明特征图如何在内存之中被空置。右侧部分表明这一节省 GPU 内存使用的想法：使用 CPU 内存作为那些特征图的临时容器。

权衡：时间 vs 空间

根据论文，vDNN（虚拟化 DNN，j2直播，也就是原文的系统设计的名称）把 AlexNet 的平均 GPU 内存使用成功降低了 91%，GoogLeNet 的降低了 95%。但你很可能已经看到，这样做的代价是训练会变慢。例如，vDNN 可以在 12GB 的 GPU 上使用 256 的批量训练 VGG-16，但是假设我们在一块拥有足够内存的 GPU 上训练同样的模型而不使用 vDNN 来优化内存使用，我们可以避免 18% 的性能损失。

当使用 cuDNN 核时，工作区大小也会出现权衡的情况。一般而言，你有的工作区越多，你可以使用的算法就越快。如果有兴趣请查阅 cuDNN 库的参考。在后面的整个讨论中我们都将会看到有关时间空间的这一权衡。

优化策略：在前向过程中卸载，在后向过程中预取

你应该已经知道 vDNN 是如何在正向过程中优化内存分配的。基本的策略是在生成特征图后将其从 GPU 上卸下传给 CPU，当它将在反向过程中被重新使用时再从 CPU 预取回 GPU 内存。这个存储空间可被释放以作他用。这样做的一个风险是如果网络拓扑是非线性的，特征图的一个张量可能被应用于数个层，从而导致它们不能被立刻卸载。当然，这个问题可以通过简单的完善优化策略来解决。