深度学习框架的内存优化机制

这篇博文简单介绍下深度学习框架的内存优化方式，主要参考资料1，也就是MXNet的官方文档。内存优化方式主要包括：in-place operation和memory sharing两种，二者在本质上没有太大差别，同时这两种也是MxNet主要采用的内存优化方式，下面的截图均来自参考资料1，接下来基于计算图来讲解。

1、in-place operation
in-place operation的示意图如下。这里有3层sigmoid，每个sigmoid的结果都放在同一个内存中，这是因为在得到第2个sigmoid的结果C后，第1个sigmoid的结果B就没有用了，所以可以直接覆盖。

但是这种in-place operation需要注意：只有当被覆盖的变量没有用了，才能覆盖。不能覆盖的情况如下图。这里当得到C后，因为B还要被另一个节点所依赖，所以C不能用B的内存。

2、memory sharing
memory sharing的示意图如下。这里当要进行第二个卷积层的时候，因为第一个卷积层的结果B已经没有用了，所以就回收第一个卷积层结果B所放的内存，用来存放第二个卷积层的结果E。

那么memory sharing有什么要求呢？原文是这么说的：The principle is to allow memory sharing only between variables whose lifetimes don’t overlap. 也就是说只要某个变量在后续的网络中用不到了，那么这个变量的内存就可以被重新利用。
这里的life time指的是一个变量从开始计算到最后一次被用到的时间。比如fc1的life time在计算得到act1后就结束了。

那么这个life time要怎么在代码中体现呢？或者说可以通过哪些方式判断某个变量的内存可以回收并重新利用呢？可以看下面这个模拟图，这里用一个计数器来记录还有多少个依赖于该节点并且还未进行计算的节点数目，只有当这个数目为0，该块内存才可以回收。右上角的矩形框表示回收到的可以被重新利用的内存，这样每次我们需要新的内存的时候既可以选择重新开辟一个内存空间，也可以选择用回收到的内存。

这里在计算B,C,E,F的时候都用到了memory sharing，在计算G的时候用到了in-place operation。

要注意的是：在上面的模拟过程中并没有真正分配内存，而是记录下每个节点需要多少内存，最后一次性申请所需的内存。

但是在进行memory sharing的时候容易遇到一个问题：并行化。想要很好地结合并行计算和内存共享是比较难的，可以看看下面这个示意图。

这是一个相同的网络，左右两个图是不同的内存共享设计图。如果不进行并行计算，那么这两个网络从A[1]到A[8]都是可以走通的。但是如果要并行计算，那么显然右边图是可以并行的，可以按网络结构分两个线程同时计算；而左边图不行，比如A[2]和A[5]不能同时计算得到，因为二者的内存是共享的。这个对比告诉我们：不要在可并行的节点之间进行内存共享。

那么要怎么判断哪些节点可以并行计算，哪些节点不能呢？可以采用类似下图这种启发式的方法来判断。首先找到网络中的最长路径，然后把这个路径用一种颜色表示，然后继续找剩下的最长的路径，再用另一种颜色表示。因为相同颜色路径的变量都有依赖关系，所以不能被并行，因此相同颜色路径的变量可以利用内存共享机制对其进行内存分配。

参考资料1：http://mxnet.io/architecture/note_memory.html
参考资料2：https://github.com/apache/incubator-mxnet/issues/797