深度学习模型的训练时内存次线性优化

深度学习模型越来越大，大公司给出的方案是分布式计算，比如Google的Mesh-Tensorflow和Gshard。但是成本对于个人开发者而言却不太友好。今天则给大家介绍一篇陈天奇大佬的工作，以计算换内存，大大提高了一个GPU上可以容纳的模型尺寸。

Overall

深度学习中的模型大部分都是层状的，因此，在正向计算和梯度计算的时候，需要我们把每一层的输出结果都保存下来，当模型层次比较深的时候，占用的内存非常大。这个时候，一个朴素的方法就是隔几层存一层，然后没有存的部分如果需要，可以实时再进行计算。

朴素的算法如下图所示。

接下来分析一下时间复杂度。假设模型一共有n层，然后分成多个segment，每个segment是k层，那么一共有n/k个segment，每个segment的输出都被保存下来。那么需要O(n/k)的复杂度，在计算一个segment的时候，需要的内存是O(k)。所以总体的时间复杂度就是O(n/k) + O(k)，这个复杂度的最小值是k=sqrt(n)的时候，所以这个方法的复杂度就是O(sqrt(n))。

但是朴素的方法有一些问题：

• 每次需要开发自己对计算图进行切分，比较麻烦且不容易找到最优的图切分方法。
• 因为是自行实现，所以很难和其他的框架自带的内存优化方法联合使用。

那么，该如何继续优化呢？且看下文。

已有的内存优化方法

我们先介绍一些已有的内存优化方法，即Inplace operation和Memory sharing。

• Inplace operation: 顾名思义，就是原地操作，如果一块内存不再需要了，且下一个操作是element-wise的，可以做原地操作覆盖内存，无需新内存。
• Memory sharing: 如果两个数据使用的内存大小一样，且有一个数据参与完计算后不再需要了，那么后一个数据可以覆盖前一个数据。

下图展示了Inplace operation和Memory sharing的例子，为了方便查看，使用的是两层神经网络。图上最右可以看到signoid可以用Inplace operation，softmax的后向计算和fc的后向计算的内存可以共享。

那么，如何找到哪些操作是inplace操作，哪些又是可以共享的内存呢？这就需要对计算图进行分析了，如下图所示，计算完后input不再需要的可以做inplace操作，生命周期不重叠的可以做共享。

优化

那么，朴素的分割方法该如何优化呢？

这里设计了一个函数m将神经网络中的每一层映射到一个非负数，其含义是这一层需要被重复计算多少次。如果m(v)=0代表这一层可以保留，m(v)>0代表这一层需要被重新计算。

所以，当所有的m(v)都等于0，那么计算图退化到正常，即所有的中间值都需要保留，对于上面那个朴素的算法，对于要保留的层次m(v)=0，其他的则是m(v)=1。

注意，m(v)有大于1的时候，此时的情况是递归分解的时候，下面会讲。

假设m已经存在，那么利用m的值对计算图进行重新安排，算法如下：

算法的效果就是找出那些需要重复计算节点。如下图所示：

有了这个算法后，如何计算m呢？论文中给出了一个贪心算法，即给定一个内存的budget，超过这个budget就保存那一层。

递归

当对计算图中的每一层进行了在内存中是保留还是丢弃的决定以后，其实就是把模型分块了。对于每一块的层次，可以递归的使用这个方法来进一步节省内存。

因为有递归的存在，空间复杂度的计算公式就变成：g(n)=k+g(n/(k+1))，从而g(n)=klogk+1(n)。

实验

使用了上述技术以后，对于1000层的resnet，可以把内存占用从48G降低到7G。对于Lstm来说，内存占用也变成了原来的四分之一。

当然，时间上，要增加30%。

其他相关知识点

• Theano和Tensorflow的内存管理使用的是GC，而MxNet则是静态分配内存。
• 很多框架都关注计算图构建完成后优化计算，但很少有框架关注计算和内存的权衡。
• 在系统研究领域，内存和计算的权衡被研究的挺多，丢弃中间层的结果的技术在自动推导图上被称作gradient checkpointing。

参考文献

• [1]. Chen, Tianqi, et al. 'Training deep nets with sublinear memory cost.' arXiv preprint arXiv:1604.06174 (2016).