MXNet整体结构预览及算子原理介绍

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 const& const_vars,
                    std::vector const& mutate_vars) = 0;

• exec_fun: 加入的函数
• exec_ctx: 函数执行的上下文环境
• const_vars: 函数读取的变量
• mutate_vars: 函数修改的变量

如果两个函数需要修改相同的变量，那么这两个函数会按照push的顺序执行

Function

using Fn = std::function;
struct RunContext {
    // stream pointer which could be safely cast to
    // cudaStream_t* type
    void *stream;
};

所有函数在引擎内部线程中执行。因此IO、UI、Web Service等任务不应该按上述方式加入引擎，因为这些耗时任务为消耗引擎内部线程计算资源，降低吞吐率。为此，MXNet提供了一下异步函数类型：

using Callback = std::function;
using AsyncFn = std::function;

在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 const& const_vars,
                                  std::vector 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 &in_data,
                    const std::vector &req,
                    const std::vector &out_data,
                    const std::vector &aux_states) = 0;
struct OpContext {
    int is_train;
    RunContext run_ctx;
    std::vector 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 &out_grad,
                    const std::vector &in_data,
                    const std::vector &out_data,
                    const std::vector &req,
                    const std::vector &in_grad,
                    const std::vector &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

    virtual bool InferShape(std::vector *in_shape,
                            std::vector *out_shape,
                            std::vector *aux_shape) const = 0;
  

  一是告诉系统输入和输出张量的尺寸，以便系统在实际计算(Forward,Backward)前分配空间。二是做shape检查用，如果不能成功推断出shape或者shape不匹配，则返回false。其中in_shape是系统根据上个算子的out_shape指定。
• Request Resources
  对于需要特定工作空间的运算来说，如果系统能够掌握这些信息，就可以进行优化，比如cudnnConvolutionForward操作。为此MXNet提供如下两个接口实现这个目的：

    virtual std::vector ForwardResource(
                const std::vector &in_shape) const;
    virtual std::vector BackwardResource(
                const std::vector &in_shape) const;
  

  目前仅包含两种类型的Resource.

  struct ResourceRequest {
        enum Type {
        kRandom,  // get a mshadow::Random object
        kTempSpace,  // request temporary space
        };
        Type type;
    };
  

  如果ForwardResource 或 BackwardResource返回非空的向量，那么在Forward和Backward运算时，可以通过ctx参数访问对应的资源。

  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算法的依赖，不被依赖的变量能够被安全释放。

  virtual std::vector DeclareBackwardDependency(
        const std::vector &out_grad,
        const std::vector &in_data,
        const std::vector &out_data) const;
  

  比如对于全连接操作，返现传播时并不需要out_data数据，而Pooling操作需要out_data数据，那么他们的反向计算依赖，可以如下定义：

  std::vector FullyConnectedProperty::DeclareBackwardDependency(
        const std::vector &out_grad,
        const std::vector &in_data,
        const std::vector &out_data) const {
        return {out_grad[0], in_data[0]};  // NOTE: out_data[0] is NOT included
    }
  std::vector PoolingProperty::DeclareBackwardDependency(
        const std::vector &out_grad,
        const std::vector &in_data,
        const std::vector &out_data) const {
        return {out_grad[0], in_data[0], out_data[0]};
    }
  

• In place Option
  指出哪些变量可共享存储空间。

  virtual std::vector>    ElewiseOpProperty::ForwardInplaceOption(
        const std::vector &in_data,
        const std::vector &out_data) const {
        return { {in_data[0], out_data[0]} };
    }
  virtual std::vector> ElewiseOpProperty::BackwardInplaceOption(
        const std::vector &out_grad,
        const std::vector &in_data,
        const std::vector &out_data,
        const std::vector &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> &kwargs) = 0;
  // return the parameters in a key-value string map
  virtual map GetParams() const = 0;
  // return the name of arguments (for generating signature in python)
  virtual vector ListArguments() const;
  // return the name of output values
  virtual vector ListOutputs() const;
  // return the name of auxiliary states
  virtual vector 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 {
    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& 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接口类和参数类

参考：MXNet System Architecture