数据抽象
NDArray
MXNet里面的NDArray其实与MShadow里面的Tensor不具有直接的联系,它里面主要采用一个Chunk的结构来存储实际的元数据,以及对多维数组的shape描述。NDArray里面同样包含一个TBlob结构,它应该是能够与Tensor进行相互转换的,但在NDArray里只是作为一种对数据存储的包装,也就是调用NDArray::data()方法的时候返回的是一个TBlob而已,这里是一个挺奇怪的设计。对于这样一个NDArray数据结构,我们主要关注以下这些方面的设计。
数据切片
从Python绑定来看,对于数据切片的支持取决于__setitem__时候key的类型,主要是integer、py_slice以及tuple这三种情况,然后最终都是调用进NDArray::Slice()这个方法。它的实现逻辑也非常简单,无非就是计算正确的数据偏移,然后维度减小,返回一个新构造的NDArray对象。我们直接去看不支持的操作,会发现它实际上没有实现带有step的切片,因为这样的操作一定会带来显式内存拷贝,看起来是没太有必要。
另一方面,对于更加复杂的选取矩阵子块的切片,这里的实现方式不是利用leading dimension,而是直接调用相应的slice操作符去进行处理,这里的实际代码实现就在MShadow里面了。
表达式执行
调度与执行
Graph Executor
Symbol Bind
如果一个抽象计算图要被执行,那么我们首先需要为它绑定执行所需要的种种信息。这个过程在MXNet中称为Symbol Bind,是通过调用Executor::Bind方法来实现的。这里我们观察一下这个函数的声明:
:::C++
static Executor *Bind(nnvm::Symbol symbol,
const Context& default_ctx,
const std::map<std::string, Context>& group2ctx,
const std::vector<NDArray> &in_args,
const std::vector<NDArray> &arg_grad_store,
const std::vector<OpReqType> &grad_req_type,
const std::vector<NDArray> &aux_states,
Executor* shared_exec = NULL);
从声明中我们大致可以猜测Bind到底是怎样的一个过程。一个计算图本质上可以理解为具有多个输入多个输出的一个过程。如果要执行一个计算图,我们首先要完成自动求导,使得图的整体结构完备。然后我们为所有输入和输出指定内存空间,使得计算图能够完成一次计算。
这里的“一次计算”指的是拓扑序遍历一遍图,执行所有节点的计算。而需要注意的是,建图时的内存优化过程使得并不是图上所有节点的内存的生存周期都能保留到计算结束,很多内存是被重用了。但是显然我们需要保留所有的权重数据,因此需要为这些数据指定显式的storage。另外,有些基本算子需要额外的辅助空间,例如BatchNorm,这些空间也需要被分配,因为它们可能需要被导出。
Context Assignment
当我们有了自动求导过后产生的计算图以后,我们能够确定计算的执行流程,梯度是如何产生的,但是还缺少实际执行所需要的一些关键信息:所有Operator的计算分别在哪里执行?
对于最简单的场景,即我们只在CPU端或者只在GPU端进行运算,那么显然这里不需要做些什么,只需要检查一下数据的一致性:输入计算图的数据NDArray所处的设备,是否与计算所执行的设备一致;以及存储梯度的数据容器(同样是NDArray)是否与所有计算所处的设备保持一致。
如果是跨设备/跨节点训练,此时我们就需要首先指定所有已知的Operator在哪个设备上执行,然后正确地调用NNVM的PlanDevice来规划设备,自动地插入数据复制节点(Copy Node)。
Execution Engine
有了经过自动求导和优化过后的计算图,此时我们就有了足够的信息用于指导我们完成整个计算流程,接下来的问题是,我们如何高效地执行这个计算图?很显然,计算图上是有并行性的,不存在数据依赖的两条路径可以直接并行执行,那么理想的方式显然是利用某种线程池机制,将所有Operator的执行过程bind成一个closure,没有数据依赖的两个closure可以一起扔进线程池,从而实现在可以并行的地方用多个线程执行的效果。那么这样的一个线程池可以用状态机的思想来实现,在每个closure运行完成之后的callback中去维护和更新状态,从而得到可以执行的下一个closure。