深度学习框架设计序列-mxnet NNVM计算图抽象机制

转自：https://blog.finaltheory.me/research/MXNet-NNVM.html 

Hint: 建议配合NNVM以及MXNet源码享用，风味最佳。

设计概要

图的抽象

简而言之就是用Symbol同时用来抽象计算图中的Operator和Operand：

• Variable Symbol
• Functor Symbol(AtomicSymbol), Callable语义

准确地讲，Symbol这个抽象概念表示一种具有多个输出的符号。Symbol对象仅仅包含一个vector<NodeEntry> outputs成员，用来记录在图中它的所有输出节点；而每一个Node记录的则是自己的属性和所有的输入inputs。所以说真正负责建图的对象是Node，它能够看到自己的所有输入；然后在计算图之上所定义的抽象概念，例如Operand和Operator使用Symbol来表示。

在图的表示中，可以用Operator来将一些Operand给compose为新的一个Operand，也就是图中的节点。最后我们拿到一个Symbol，从它就可以回溯出整个图，转换为Graph对象来完成建图操作。

不同的Operator会有自己的属性，常用的属性定义在op_attr_types.h里面，基本上就是不同的类型声明。

Python接口

这里的实现非常漂亮，首先C API定义的导出接口就很少，Python代码中仅仅定义最核心的Symbol类、运算符重载以及必要的ctypes胶水代码。至于如何用Operator来建图，则是在C++代码中通过静态定义注册进去，然后在import nnvm的时候动态注册进Python。具体实现里面有更多的细节。

实现细节

自动求导

基本原理

自动求导的原理比较简单，注意所有变量都是Tensor。考虑操作符Operator的基本性质与函数稍有区别，是从多个输入映射到多个输出。所以设当前节点的操作符接受m个输入，返回n个输出，则形式化表示为：

Op(x1,x2,...,xm)=[f1(x1,...,xm),...,fn(x1,...,xm)]Op(x1,x2,...,xm)=[f1(x1,...,xm),...,fn(x1,...,xm)]

其中f1,...,fnf1,...,fn为不同输出所对应的运算过程（函数）。

为每个操作符定义Gradient运算过程，输入为误差yy对当前节点所有输出的导数（一个list），输出为误差对当前节点所有输入的导数（同样是一个list）。输入形式化表示为[∂y∂f1(...),∂y∂f2(...),...,∂y∂fn(...)][∂y∂f1(...),∂y∂f2(...),...,∂y∂fn(...)]，那么输出就应该是：

[∑i=1n∂y∂fi∂fi∂x1,∑i=1n∂y∂fi∂fi∂x2,...,∑i=1n∂y∂fi∂fi∂xm][∑i=1n∂y∂fi∂fi∂x1,∑i=1n∂y∂fi∂fi∂x2,...,∑i=1n∂y∂fi∂fi∂xm]

另外，∂y∂fi(...)∂y∂fi(...)是把从yy到操作符Op的某个输出fi(...)fi(...)的所有偏导路径都求和之后，累加得出的结果。这也就意味着所有操作符都需要定义累加运算的属性，如果没有的话，在MXNet中默认会尝试调用__ewise_sum__进行累加。注意以上这些运算过程都是符号运算而非数值运算，也就是运算会产生图上新的节点。

接下来我们就直接按照逆拓扑序遍历网络，把当前节点的所有输出的导数的聚合结果（一个节点）[∂y∂fi(...)][∂y∂fi(...)]求出来，然后用操作符的Gradient方法求出对所有输入的导数[∂y∂xi][∂y∂xi]，再不断反向传播即可。逆拓扑序可以保证遍历到当前节点的时候，它的所有输出节点，也就是反向传播时的所有输入节点，都已经求导完成了。

在神经网络中，我们做梯度下降需要误差yy对权重矩阵WW的导数，所以在自动求导的过程中，我们只需要将关心的这些权重节点的导数节点输出即可。

MXNet实现细节

下图为一个简单的两层MLP的Forward过程计算图：

那么在反向求导的过程就是这个样子：

1. 求出softmax对fc3和softmax_label的导数softmax_backward.
2. 此时已知yy对fc3节点的所有输出的导数是[softmax_backward]，那么就可以计算yy对fc3的所有输入[relu2, fc3_weight, fc3_bias]的导数。
3. 以此类推，逆拓扑序完成整个求导过程……

这里我们注意到，softmax具有两个shape明显不同的输入节点，但却只有一个导数节点softmax_backward；同样地fc3有三个不同的输入，但却也只有一个导数节点fc3_backward，其他节点的情况都与此类似。这里实际上是MXNet重构计算图设计的历史遗留问题。因为绝大多数旧操作符并不是按照上述设计来实现的——他们并不包含一个Gradient方法，能够根据节点的输入直接按照符号运算求出导数节点。所以为了兼容，MXNet注册操作符的时候在NNVM中为这些旧版本的操作符注册了加前缀_backward_的反向操作符。在对计算图进行求导的时候，会直接生成一个Operator为_backward_版本的节点，并且让操作符声明自己在反向求导过程中所依赖的数据。然后这个_backward_节点在实际执行计算图的时候，会按照反向的行为进行计算，得到正确的导数结果。

Type/Shape Inference

通过自动求导构建计算图以后，我们就能以输入Tensor的shape，推导出图中所有节点的shape，从而用于后续的内存分配等过程。这一步要求各个Operator实现一个InferShape方法，其输入是当前节点指针以及该操作符的所有输入数据的shape，返回值是该操作符所有输出数据的shape。对于每个操作符来说，实现这样的功能是很简单的：比如对于element-wise加法，那么输出数据的shape就等于输入数据的shape；对于矩阵乘法则需要考虑是否转置之类的情况。

由于上文提到的MXNet的历史遗留问题，对于反向过程要更麻烦一些，需要针对性的特殊处理。这里面用到一个“控制流依赖”(control flow dependencies)的关系来进行反向过程的类型推导。至于具体怎么实现的……这代码我是不想继续纠结了。

内存分配

首先贴上架构设计文档：Optimizing Memory Consumption in Deep Learning。

为计算图的所有节点分配内存的问题可以抽象为：给定一个内存块Request/Free的操作序列，要求满足所有分配的需求，并且使得总的分配内存数量最少。所以这是个NP问题么？

内存分配器会有一个参数match_range_，用来表示在[size/match_range_, size*match_range_]的范围内来寻找内存块。这里面的trick在于先试图分配大的内存块，然后找不到的话再试图分配小的内存块。当然这里并不是真的在分配内存，而只是预先规划我要分配怎样的内存，如果找到小的内存块，肯定不满足我们的需求，我们此时把它放大到我们想要的size即可，我们现在是在记录需求，反正最终运行的时候才会真的分配内存。

接下来分析实现细节。整体上分为初始化阶段和按照拓扑序遍历计算图的阶段。

初始化阶段

1. 内存分配阶段要依赖于Shape和Type的Inference，这是显然的，不然分配个毛啊。在注册这个Pass的时候，会指定这种依赖关系。
2. 然后要计算所有非Variable节点的出度，作为refcount；有些操作符具有FIgnoreInputs属性，并不需要输入数据（只要shape），比如zeroslike这样的操作符，所以遍历的时候不要算这部分的引用计数。
3. 输出节点要额外加一个引用计数（出度+1），保证在计算图执行到结束的时候也不会回收这些内存。这一点很重要，我就踩过坑。

拓扑序遍历阶段

这一阶段直接是一个for循环，以拓扑序遍历整个计算图，循环体内所做的事情如下：

1. 首先是检查是否能做in-place运算优化。Operator可以设置自己支持inplace操作来显式优化内存分配，所以内存分配的时候是先处理能够inplace的情况，然后再操作正常的内存分配。另外inplace优化实际上可能是一对多的关系，就是说运算符可以指定一个输入节点的内存可能被复用给多个输出节点，因为可能有的输出节点只需要shape信息，不需要数据本身，根本不用给他分配空间。最后，inplace优化需要满足一个挺复杂的条件：
   • 输入节点只对应一个输出（出度为1）
   • 输出节点有被其他节点引用（否则就不需要为它分配内存，因为根本不用算它）
   • 输出节点尚未分配内存
   • 输入节点已分配内存（拓扑序遍历的话，这个条件应该是默认满足的）
   • 数据类型、大小匹配
2. 接下来就开始遍历当前节点的所有输出了，把所有还没分配内存的节点记录下来排个序，从小到大依次向内存分配器请求内存即可。
3. 然后我们就可以更新引用计数了：把所有输入节点（排除FIgnoreInputs节点）的refcount - 1，如果refcount == 0，就可以释放这个节点的内存。另外这时会遇到有些节点出度本来就是零，这是因为inplace优化导致的，跳过就行了。
4. 最后我们还需要遍历一遍输出节点，把那些出度为零的节点的内存释放掉，标记为不需要分配内存，因为他们根本不被用到，对用户来说处于“不可见状态”。

设备分配

这个pass的功能就是在计算的时候设置不同的节点/子图在哪个设备上进行计算。如果出现了跨设备的数据依赖，就增加一个设备之间数据拷贝的节点。显然，这样的设备分配策略会改变计算图的结构，所以这里采用经典的持久化数据结构，即只增加不修改，从而基于原计算图得到新的计算图。它的策略看起来相对简单：

1. 初始状态下，为所有节点赋予对应的设备编号device_id，默认为-1（invalid）；
2. 然后开始拓扑序遍历计算图，每个节点可以包含一个属性(string)，指明这个节点在计算的时候属于哪个分组（group）；同时计算图自己也会有一个设备分配的映射关系，是从group到计算设备（device_id）的映射。这样如果一个节点有分组信息，就可以直接根据计算图的映射来找到它所对应的设备了。如果节点木有设备分组属性的话，就直接把当前节点的输入节点的设备作为当前节点的运算设备，这是显然的。
3. 我们还需要逆拓扑序遍历一遍。猜测这是因为正拓扑序遍历能够保证所有前向运算分配到合适的设备，但是逆拓扑序才能保证给反向运算分配合适的设备。此时所有的节点就都具有自己的device_id了。
4. 接下来开始为所有跨设备的运算操作之间插入数据复制的Operator。这是一个相对复杂的过程：
   • 首先要进行一个合法性检查，因为NNVM的设计允许Operator去就地mutate它的inputs，但如果当前节点与它的输入节点并不在一个设备上面，就不能这样做，因为这样的就地mutate不能通过插入一个copy节点来实现。
   • 然后检查对于当前节点是否需要变更图结构。假设我们知道经过所有操作以后，我们生成了一个新图，那么显然在新图中的每个节点和原图存在一一对应的关系。这里采用了一个new_node_map的哈希表，来表示某个节点在新图中所对应的节点。
   • 如果当前节点的输入inputs中有节点需要被映射到新的节点，那么我们就也需要为当前节点创建一个新的映射节点，然后将那个对应的输入节点指向它所对应的新节点。
   • 并且，如果当前节点的某个输入所对应的device_id与其自身的device_id不同，那么我们就插入一个新的copy节点，用于在执行的时候将数据正确地从不同设备之间进行拷贝。然后这个新的节点就可以被加入new_node_map中。
5. 最后我们返回一个新图，方法就是把旧图中的outputs节点替换成new_node_map中的对应节点。

一个简单的例子，考虑如下代码：

import mxnet as mx
a = mx.sym.Variable('a')
b = mx.sym.Variable('b')
c = mx.sym.Variable('c')
with mx.AttrScope(ctx_group='dev1'):
    net = a * b
with mx.AttrScope(ctx_group='dev2'):
    net = net + c
e = net.simple_bind(mx.cpu(), a=(10, 10), grad_req='write',
                    group2ctx={'dev1': mx.cpu(), 'dev2': mx.gpu()})
mx.viz.plot_network(net)

也就是说，默认整个net运行在CPU上面，指定net = a + b运行在CPU上；但是指定其中一个步骤net = net + c运行在GPU上面。

则生成的计算图如下：

经过了Plan Device Pass之后，该计算图被修改为如下所示：

可以看到，为了满足中间计算过程的跨设备要求，c以及a * b的运算结果被显式拷贝操作符复制到GPU设备上面，进行加法运算以及反向加法求导之后，又被拷贝回CPU端来继续进行乘法求导。

代码风格&槽点

• 一句话概括整体风格就是有种用Python思想写C++的感觉。
• Operator的注册机制（Registry）有点滥用全局状态，而且有些优化trick显得意义不大，宏接口设计的倒是比较漂亮。
• 尽管template到处飞，但是静态类型整体上用的不是特别好，到处出现的Op::GetAttr<FGradient>()里面那个参数字符串看得人难受。
• 计算图优化的部分采用一个pass一个编译单元的结构，代码质量一般，还是觉得在建图的过程中做了一些用处不大的优化trick，可读性不是很好。
• Python接口与ctypes部分写得很赞，与NNVM_REGISTER_OP的协作非常漂亮，使得运算符注册能够在import nnvm的时候自动搞定。就是……有点绕