MXNet Architecture
Focus on 3 areas: abstraction, optimization, and trade-offs between efficiency and flexibility. The document is an overview of the complete MXNet system.
MXNet System Architecture
MXNet System Architecture: major modules and components of the MXNet system and their interaction.
- Runtime Dependency Engine
根据读写依赖调度和执行计算操作。 - Storage Allocator
分配和回收CPU与GPU存储空间。 - Resource Manager
管理全局资源，比如随机数生成器和temporal space. - Operator
定义的计算算子。 - NDArray
动态异步多维数组，支持高效计算图执行和优化。 - Symbolic Execution
静态符号图执行器，提供高效的符号图执行和优化。 - SimpleOP
以统一的方式扩展NDArray运算符和符号运算符。 - Symbol Construction
计算图构造器。也可通过net configuration完成。 - KVStore
高效的键值参数同步接口。 - Data Loading(IO)
高效的分布式数据加载和扩充。
Execution Engine
执行引擎不仅仅用于深度学习，它也能解决其它领域的问题。它的设计目的是：按照依赖关系执行一堆函数。如果函数有依赖关系，引擎能够保证函数根据依赖关系顺序执行，如果没有依赖关系函数能够并行执行。
Interface
virtual void PushSync(Fn exec_fun, Context exec_ctx,
std::vector<VarHandle> const& const_vars,
std::vector<VarHandle> const& mutate_vars) = 0;
- exec_fun: 加入的函数
- exec_ctx: 函数执行的上下文环境
- const_vars: 函数读取的变量
- mutate_vars: 函数修改的变量
如果两个函数需要修改相同的变量，那么这两个函数会按照push的顺序执行
Function
using Fn = std::function<void(RunContext)>;
struct RunContext {
// stream pointer which could be safely cast to
// cudaStream_t* type
void *stream;
};
所有函数在引擎内部线程中执行。因此IO、UI、Web Service等任务不应该按上述方式加入引擎，因为这些耗时任务为消耗引擎内部线程计算资源，降低吞吐率。为此，MXNet提供了一下异步函数类型：
using Callback = std::function<void()>;
using AsyncFn = std::function<void(RunContext, Callback)>;
在AsyncFn中我们能传入自定义线程执行耗时任务，引擎直到Callback被调用时才会认为AsyncFn已经执行结束。
Context
Context定义了函数在什么环境下执行，比如支持函数在CPU或GPU上执行，函数执行时使用的设备id是什么等信息。RunContext与Context是不同的，RunContext包含只有执行时才能确定的信息，比如函数要使用哪个类型的数据流。
VarHandle
用来定义函数依赖。创建变量使用NewVar接口, 删除变量使用PushDelete接口.
Push and Wait
所有Push操作都是异步的，而且不保证线程安全。通过Wait接口可以实现序列话。
- 使用AsyncFn能够监控特定函数执行完成。
- 使用WaitForVar监控变量var已被所有函数使用(读或写)完毕。
- 使用WaitForALl能够监控所有已Push的函数执行完毕。
Save Object Creation Cost
对于长期push多个函数时，拷贝函数lambdas和创建变量都较为耗时，这种情况下，我们可以预创建OprHandle来节省时间。
virtual OprHandle NewOperator(AsyncFn fn,
std::vector<VarHandle> const& const_vars,
std::vector<VarHandle> const& mutate_vars) = 0;
virtual void Push(OprHandle op, Context exec_ctx) = 0;
当Push操作完成后，就可以调用DeleteOperator(OprHandle op)删除预创建的函数句柄。
Operators
在MXNet中一个算子是一个class，包括计算逻辑和优化辅助信息两个部分。
MXNet算子接口一些功能：
- 通过制定in-place更新，减少内存分配。
- 隐藏一些Python调用的内部参数，保持算子简洁。
- 通过定义输入张量和输出张量的关系，实现张量shape检查。
- 从系统中申请额外的时空资源。（如，calling cudnn routines)。
Operator
算子逻辑实现接口类，主要包括Forward和Backward两个接口。
virtual void Forward(const OpContext &ctx,
const std::vector<TBlob> &in_data,
const std::vector<OpReqType> &req,
const std::vector<TBlob> &out_data,
const std::vector<TBlob> &aux_states) = 0;
struct OpContext {
int is_train;
RunContext run_ctx;
std::vector<Resource> requested;
}
enum OpReqType {
kNullOp, // no operation, do not write anything
kWriteTo, // write gradient to provided space
kWriteInplace, // perform an in-place write
kAddTo // add to the provided space
};
- ctx: 操作是训练还是推断阶段，操作运行环境，所需的资源。
- in_data, out_data: 输入与输出张量。
- req: 指出每个out_data的写入策略。req[i]表示out_data[i]的写入策略。kWriteTo表示对应的out_data是raw memory block, 可以直接写入。kAddTo表示累加计算结果和内存中现有值，实现运算"+="。
- aux_states: 辅助张量运算。目前还未使用。
virtual void Backward(const OpContext &ctx,
const std::vector<TBlob> &out_grad,
const std::vector<TBlob> &in_data,
const std::vector<TBlob> &out_data,
const std::vector<OpReqType> &req,
const std::vector<TBlob> &in_grad,
const std::vector<TBlob> &aux_states);
与Forward的类似的设计原则。只是这个借口中out_grad, in_data, out_data是已知的，要求的计算结果放在in_grad中。有些算子在计算输入梯度时，并不需要全部out_grad, in_data和out_data三个变量，此时使用OperatorProperty中的DeclareBackwardDependency接口声明只需要那几个变量。不需要的变量引擎会根据需要自动释放空间。
Operator Property
分离算法语义与实现，达到一个语义算子支持不同实现的目的。
语义接口类：OperatorProperty，实现接口类：Operator。
算子语义接口类OperatorProperty主要组成部分如下：
- InferShape
一是告诉系统输入和输出张量的尺寸，以便系统在实际计算(Forward,Backward)前分配空间。二是做shape检查用，如果不能成功推断出shape或者shape不匹配，则返回false。其中in_shape是系统根据上个算子的out_shape指定。virtual bool InferShape(std::vector<TShape> *in_shape, std::vector<TShape> *out_shape, std::vector<TShape> *aux_shape) const = 0; - Request Resources
对于需要特定工作空间的运算来说，如果系统能够掌握这些信息，就可以进行优化，比如cudnnConvolutionForward操作。为此MXNet提供如下两个接口实现这个目的：
目前仅包含两种类型的Resource.virtual std::vector<ResourceRequest> ForwardResource( const std::vector<TShape> &in_shape) const; virtual std::vector<ResourceRequest> BackwardResource( const std::vector<TShape> &in_shape) const;
如果ForwardResource 或 BackwardResource返回非空的向量，那么在Forward和Backward运算时，可以通过ctx参数访问对应的资源。struct ResourceRequest { enum Type { kRandom, // get a mshadow::Random object kTempSpace, // request temporary space }; Type type; };auto tmp_space_res = ctx.requested[kTempSpace].get_space(some_shape, some_stream); auto rand_res = ctx.requested[kRandom].get_random(some_stream); - Backward dependency
定义Backward算法的依赖，不被依赖的变量能够被安全释放。
比如对于全连接操作，返现传播时并不需要out_data数据，而Pooling操作需要out_data数据，那么他们的反向计算依赖，可以如下定义：virtual std::vector<int> DeclareBackwardDependency( const std::vector<int> &out_grad, const std::vector<int> &in_data, const std::vector<int> &out_data) const;std::vector<int> FullyConnectedProperty::DeclareBackwardDependency( const std::vector<int> &out_grad, const std::vector<int> &in_data, const std::vector<int> &out_data) const { return {out_grad[0], in_data[0]}; // NOTE: out_data[0] is NOT included } std::vector<int> PoolingProperty::DeclareBackwardDependency( const std::vector<int> &out_grad, const std::vector<int> &in_data, const std::vector<int> &out_data) const { return {out_grad[0], in_data[0], out_data[0]}; } - In place Option
指出哪些变量可共享存储空间。virtual std::vector<std::pair<int, void*>> ElewiseOpProperty::ForwardInplaceOption( const std::vector<int> &in_data, const std::vector<void*> &out_data) const { return { {in_data[0], out_data[0]} }; } virtual std::vector<std::pair<int, void*>> ElewiseOpProperty::BackwardInplaceOption( const std::vector<int> &out_grad, const std::vector<int> &in_data, const std::vector<int> &out_data, const std::vector<void*> &in_grad) const { return { {out_grad[0], in_grad[0]} } } - Expose Operator to Python
为了将算子暴露给Python语言，其它需要实现的接口有：// initial the property class from a list of key-value string pairs virtual void Init(const vector<pair<string, string>> &kwargs) = 0; // return the parameters in a key-value string map virtual map<string, string> GetParams() const = 0; // return the name of arguments (for generating signature in python) virtual vector<string> ListArguments() const; // return the name of output values virtual vector<string> ListOutputs() const; // return the name of auxiliary states virtual vector<string> ListAuxiliaryStates() const; // return the number of output values virtual int NumOutputs() const; // return the number of visible outputs virtual int NumVisibleOutputs() const;
Create an Operator from the Operator Property
通过OperatorProperty接口类创建具体算子。以具有参数的算子为例，主要代码如下：
struct ConvolutionParam : public dmlc::Parameter<ConvolutionParam> {
TShape kernel, stride, pad;
uint32_t num_filter, num_group, workspace;
bool no_bias;
};
class ConvolutionOp {
public:
ConvolutionOp(ConvolutionParam p): param_(p) {}
void Forward( ... ) { ... }
void Backward( ... ) { ... }
private:
ConvolutionParam param_;
};
class ConvolutionOpProperty : public OperatorProperty {
public:
void Init(const vector<pair<string, string>& kwargs) {
// initialize param_ using kwargs
}
Operator* CreateOperator(Context ctx) const {
return new ConvolutionOp(param_);
}
private:
ConvolutionParam param_;
};
定义的算子参数，算子语义和算子实现，要注册到MXNet。
DMLC_REGISTER_PARAMETER(ConvolutionParam);
MXNET_REGISTER_OP_PROPERTY(Convolution, ConvolutionOpProperty);
MXNET_REGISTER_OP_PROPERTY的第一个参数是名称字符串，第二个参数要注册的OperatorProperty类名称。
interface Summary
创建一个算子的一般步骤为：
- 实现Operator接口类，编写计算逻辑代码(Forward, Backword)
- 实现OperatorProperty接口类，完成以下功能
- 向Operator接口类传入所需要的参数（可通过Init接口实现）。
- 使用CreateOperator接口创建Operator接口类的实例。
- 实现算子接口描述，比如参数名称等。
- 实现InferShape接口，指出输出张量的shape。
- 【可选】如果需要额外资源，实现ForwardResource和BackwardResource接口。
- 【可选】如果Backward计算不需要所有输入输出变量，实现DeclareBackwardDependency接口指出依赖关系。
- 【可选】如果支持in-place更新，实现ForwardInplaceOption和BackwardInplaceOption接口。
- 注册实现的OperatorProperty接口类和参数类
