「人工智能研学社· ML系统与架构小组」第一期：如何在单块GPU上训练超大型深度学习模型

机器之心 · 公众号 · AI · 2017-04-16 11:29

正文

请到「今天看啥」查看全文

减小你的批量大小，但这可能会妨碍你的训练速度和精确度。
在多 GPU 环境下做模型并行，这是另一个复杂的事情。
缩小你的模型，如果你不情愿做出上述两个选择，或者已经尝试但效果不好。

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

观察：什么占用内存？

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

模型参数
特征图
梯度图
工作区

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

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

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

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

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

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

权衡：时间 vs 空间

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

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