正文
-
减小你的批量大小,但这可能会妨碍你的训练速度和精确度。
-
在多 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 库的参考。在后面的整个讨论中我们都将会看到有关时间空间的这一权衡。