教程|如何在单块GPU上训练大型深度模型?

机器之心原创

参与机器之心编辑部

问题：GPU 内存限制

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

在给定模型和批量大小的情况下，事实上你可以计算出训练所需的 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 内存中的数据分为 4 个部分：

• 模型参数（权重）

• 特征映射

• 梯度映射

• 工作区

前 3 个功能性容易理解。每个人都知道权重是什么。特征映射是正向过程中生成的中间结果。梯度映射是反向过程中生成的中间结果。工作区是 cuDNN 函数临时变量／矩阵的一个缓冲区。对于一些 cuDNN 函数，用户需要通过缓冲区到达作为函数参数的内核。一旦返回函数，该缓冲区将被释放。

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

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

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

这一想法是：如果 GPU 内存中的大部分数据出现空置，为什么不把它们保存在更便宜的 CPU 内存上呢？这里是一张在 AlexNet 的实例中表明什么在进行的图表。

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

权衡：时间 vs 空间

根据论文，vDNN（可视化DNN）把 AlexNet 的平均 GPU 内存使用成功降低了 91% ， GoogLeNet 的降低了 95%。你很可能已经看到，这样做的代价是训练会变慢。例如，vDNN 可以在 12GB 的 GPU 上使用 256 的批量训练 VGG-16，但是相比于带有足够内存的假设性 GPU ，这会产生 18% 的性能损失。

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

优化策略：在正向与反向的过程中

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

在后向过程中，vDNN 采用一种更具侵略性的策略，这是由于相较于特征映射，梯度映射不存在“稍后重用”的问题。因此一旦生成了相关的权值更新，它们就可以被释放 （相较于那些映射，这些权值更新是很小的）。

优化策略：内存管理器 CUDA 流

vDNN 的关键部分是 cuda 流，它管理内存分配／释放、卸载和预取。这里是一些细节：

传统的 cuda 内存分配／释放（cudaMalloc & cudaFree）是同步性 API。由于随着训练过程的进行连续地发生，同步性 API 并不够高效。

如同分配/释放操作，卸载 API 也需要是异步性的。当选择卸载特征映射的一个张量时，vDNN 的内存管理器流（ memory manager stream of vDNN）将在主机上分配一个固定的内存管理区域，并且通过 PCIe 发出一个非阻塞传输（ a non-blocking transfer）。这些特征映射张量在正向过程中为只读的，因此这个传输过程可以在计算中被安全地覆盖。只有当前层的卸载完成时，才能进行下一层的处理。

预取操作是在反向处理中从 CPU 返回到 GPU 以得到卸载的特征映射。和上面的操作类似，预取操作也需要是异步性的。显而易见，在预取和相同层的计算之间存在数据依赖，因此 vDNN 将同时异步开始当前层的计算以及前一层的预取。

成本：如何偿还为节省内存而付出的性能代价？

最显著的潜在性能损失来自于由卸载／预取引入的隐式依赖（implicit dependency）。考虑到数据传输比正向计算需要花费更长的时间。这一图表清晰地表明这种情况（图 9：卸载和预取的性能影响(Minsoo Rhu et al. 2016)）

相似情形也可能在反向过程中发生。

算的情况下，如何获得最佳性能？

如上所述，在时间和空间之间有一个权衡，并且在前的章节中我们已经看到权衡是如何工作的。想象一下你正在 12GB 的 GPU 上使用 128 的批量（这需要 14GB 内存如果没有卸载／预取）训练 VGG-16。仅使用了大约 2GB 的内存也许很浪费，因为你本可以使用更多的空间减少性能损失。因此，我们可以这种方式重新形式化这个问题：在限制内存预算的情况下，如何获得最佳性能？

配置时间-空间权衡：决定一个层是否应该被卸载/预取（offloaded/prefetched），以及应该选择哪种卷积算法。

为了取得最佳配置，我们需要为每个层决定两件事：我们是否需要卸载/预取，我们在前向/反向传播过程中该选择哪个算法（更快收敛的算法需要更大的存储空间）。

通常，靠近输入的层有较长重用度（reuse ）距离。因此，首先卸载/预取这些层比较好。现在我们并不需要为每一层决定是否使用（选择随层级数成指数增长）。我们只需要选择一层，并且更接近输入的每一层都进行卸载/预取，其余层将其张量保留在 GPU 上。

为每层决定算法也并不可行（同样选择随层级数成指数增长）。我们可以通过强制每层使用相同的算法（gemm 或fft 等算法）来简化这一问题。

现在我们有一个小型配置空间，这样我们可以尽情搜索以确定最佳配置方案。 下面的图表说明了我们的配置空间：

左上角表示存储器优化配置（卸载/预取每一层，并使用最慢的算法），右下角表示性能最优的配置（当然，真实的配置空间应该是网格，可行模型与不可行模型之间的边界应改成阶梯状，不过这一图表足以说明这一点）。

下一步即找到具有最佳性能的可行配置。 我们可以沿可行性边界搜索配置。 如果你对该搜索过程的实现细节感兴趣，请参阅原论文（搜索过程在第 3.C 节中进行了描述/https://arxiv.org/pdf/1602.08124.pdf）。