mxnet-增加新层(cpp)

重点

• op-inl.h文件中是主要代码，以内联函数和模板为主，代码可以同时支持cpu和gpu
• .cc文件中是cpu相关代码，包括模板的cpu版实例化
• .cu文件中是gpu相关代码，包括模板的gpu版实例化
• 实现一个operator包括六个步骤：注册参数，实现shape/type推测函数，forward，backward，注册operator，单元测试
• nvcc负责注册operator
• mshadow是对gpu/cpu并行代码的封装，令一份代码同时支持cpu和gpu。一般cpu使用openMP实现，gpu就是cuda实现

A Beginner’s Guide to Implementing Operators in MXNet Backend

Introduction

构造神经网络最核心的元素是operator. operator定义了输入到输出转换的数学公式.从最简单的operator,比如元素求和,到复杂如卷积,mxnet包含了丰富的operator支持当前流行的各种神经网络. mxnet中的operator具有和numpy相似的接口,比如repeat,tile等等. 重新构造轮子的一个主要原因是numpy不支持GPU运算,而mxnet需要同时支持CPU和GPU.此外在mxnet各个组件上我们都做了大量的优化,比如tensor,execution engine和computation graph等,极大的提高了速度和内存操作的效率.
本教程中我们将在c++端实现一个operator,然后在python端进行测试.

Implementation

An Operator Example

以二次曲线为例: $f(x) = ax^2 + b + c$, 我们准备实现一个名为quadratic的operator,输入tensor x,输出tensor y. x和y的尺寸一致, y的每个元素等于x对应元素输入二次函数f的结果. 变量a,b和c将作为输入参数. 在前端这个operator按如下方式调用.

x = [[1, 2], [3, 4]]
y = quadratic(data=x, a=1, b=2, c=3)
y = [[6, 11], [18, 27]]

我们需要先创建三个文件: quadratic_op-inl.h, quadratic_op.cc和quadratic_op.cu. 头文件命名包括operator名字,op和-inl三个部分,表示这是一个以内联函数实现的operator,支持CPU和GPU计算. CPU和GPU无法共享的代码分别在cc和cu文件中. 我们一般把纯tensor相关的operator(诸如tile,repeat等)放在src/operator/tensor目录中, 和神经网络相关的operator(诸如convolution,pool等)放在src/operator/nn目录中.现在的版本中很多和神经网络相关的operator,比如Convolution和Pooling放在了src/operator中, 为了更好,更清晰的文件目录组织结果,未来的版本他们将会被安置在src/operator/nn目录中.
下面我们将进行如下几个步骤:
1. 在quadratic_op-inl.h中定义参数结构体,注册参数a,b和c
2. 在quadratic_op-inl.h中定义type和shape的推测函数
3. 在quadratic_op-inl.h中定义forward和backward函数
4. 分别在quadratic_op.cc和quadratic_op.cu文件中用nnvm注册operator

Parameter Registration

首先在quadartic_op-inl.h中定义struct QuadraticParam作为参数a,b,c的占位符(译注:延迟计算). 这个struct继承自一个名为dmlc:Parameter的模板struct, 该模板的参数就是QuadraticParam. 这项技术被称为”curiously recurring template pattern”, 和虚函数类似,但是它实现了静态多态,而不是动态多态.

struct QuadraticParam : public dmlc::Parameter<QuadraticParam> {
  float a, b, c;
  DMLC_DECLARE_PARAMETER(QuadraticParam) {
    DMLC_DECLARE_FIELD(a)
      .set_default(0.0)
      .describe("Coefficient of the quadratic term in the quadratic function.");
    DMLC_DECLARE_FIELD(b)
      .set_default(0.0)
      .describe("Coefficient of the linear term in the quadratic function.");
    DMLC_DECLARE_FIELD(c)
      .set_default(0.0)
      .describe("Constant term in the quadratic function.");
  }
};

上述代码中的函数名足以解释其含义. 每个参数设置了默认值0.0, 允许用户传递0个参数调用operator. 不设置默认值的参数需要用户在运行是输入值, 另外每个参数都附带一个简短的说明,通过文档引擎在mxnet documentation web page上显示.

Attribute Inference

Attribute inference是一个根据用户提供的信息推测网络中的NDArray的属性. NDArray两个最常见的属性是type和shape. 已知输入的NDArray名字是data, 通过

output = mx.nd.quadratic(data, a=1, b=2, c=3)

调用quadratic operator. 在开始计算之前,要先通过你定义的规则,从输入的data推测type和shape,以分配内存.

推测函数的重点是实现mutual inference: 根据operator的定义,从一个参数的属性推测另一个参数的属性. 符号编程中计算图自动推测未知属性是十分有用的能力. 用户把计算图看作一个为了让数据顺利在网络中流动的而初始化的过程,包括内存分配,设备指定等.用户仅仅需要给计算图提供最少的信息,比如输入数据的shape, 计算图将自动推测出未知的属性, 建立起整个神经网络.下面是一个例子

>>> import mxnet as mx
>>> a = mx.sym.Variable('a', shape=(2, 0))
>>> b = mx.sym.Variable('b')
>>> c = mx.sym.Variable('c', shape=(0, 3))
>>> d = a * b + b * c
>>> print d.infer_shape()
([(2L, 3L), (2L, 3L), (2L, 3L)], [(2L, 3L)], [])

上面的例子显示了d.infer_shape()返回的包含三个list的tuple. 第一个list包含所有参数a,b和c的shape. 第二个包含输出d的shape, 第三个包含一些附加参数的shape(本例子中没有使用,所以是空list). 例子中仅指定了a的第一个维度和c的第二个维度. (2,0)中的0表示第二个维度的尺寸未知,(0,3)中的0含义类似. 最终符号d成功推测出所有参数以及输入符号的shape,这就是mutual inference的结果. mxnet中整个推测过程如下:
1. a和b通过逐元素乘法关联,所以a和b的shape是一致的,因此b的第一个维度尺寸是2
2. b和c也通过逐元素乘法关联,所以b和c的shape是一致的,因此b的第二个维度尺寸是3
3. 至此b的shape是已知的了,所以a和c未知的维度尺寸也是已知的了
4. d是$a*b$和$b*c$的和,所以d的shape和b一致
上述四步展示了mxnet中shape推测是如何进行的. 这个逻辑在逐元素乘法和加法的operator的shape inference函数中编码实现.
对于quadratic operator, shape inference的逻辑类似:

inline bool QuadraticOpShape(const nnvm::NodeAttrs& attrs,
                             std::vector<TShape>* in_attrs,
                             std::vector<TShape>* out_attrs) {
  CHECK_EQ(in_attrs->size(), 1U);
  CHECK_EQ(out_attrs->size(), 1U);

  SHAPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(0));
  SHAPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(0));
  return out_attrs->at(0).ndim() != 0U && out_attrs->at(0).Size() != 0U;
}

上面例子中有些需要注意的地方:
1. attrs中有用户输入的参数a,b和c.例子中因为不需要而没有使用.
2. in_attrs是一个包含了所有输入shape的向量.因为quadratic operator只有一个输入,我们利用CHECK_EQ检查向量元素个数是否正确
3. out_attrs是一个包含了所有输出shape的向量,同样因为只有一个输出,例子中用CHECK_EQ检查元素个数
4. 两次调用SHAPE_ASSING_CHECK实现mutual inference. 第一个通过输入shape推测输出shape,第二个通过输出shape推测输入shape. 如果两个shape中有任何不相等的非零值,比如(2,3)和(3,3), 这个宏将抛出异常和shape inference的错误信息
5. 函数的最后,通过检查输出shape非空,并且每个值都大于0来确保输出shape都已经确定. mxnet中,空shape表示shape未知,而shape中的0值表示对应维度上的尺寸未知.不论哪一种情况,未知的shape信息都要从另一个shape中推测出来.如果无法推测出来,函数返回false,并通知调用者shape inference失败.
6. 这个例子只是为了展示,因为mxnet已经为逐元素操作的operator提供了方便的实现,完成mutual inference. 用户可以设置n_in=1, n_out=1,即实现了上述例子中的QuadraticOpShape.

template<int n_in, int n_out>
inline bool ElemwiseShape(const nnvm::NodeAttrs& attrs,
                          std::vector<TShape> *in_attrs,
                          std::vector<TShape> *out_attrs);

type inferenc是同样的逻辑, 我们提供下面的例子供你们分析. data type中的-1表示type未知,需要通过输入/输出type推测出来.

inline bool QuadraticOpType(const nnvm::NodeAttrs& attrs,
                            std::vector<int>* in_attrs,
                            std::vector<int>* out_attrs) {
  CHECK_EQ(in_attrs->size(), 1U);
  CHECK_EQ(out_attrs->size(), 1U);

  TYPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(0));
  TYPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(0));
  return out_attrs->at(0) != -1;
}

同样mxnet为逐元素操作的operator提供了一个方便的函数,如下所示

template<int n_in, int n_out>
inline bool ElemwiseType(const nnvm::NodeAttrs& attrs,
                         std::vector<int>* in_attrs,
                         std::vector<int>* out_attrs);

Forward Function

Forward函数定义了网络前向过程中operator的行为. quadratic operator仅需逐元素的调用二次函数. mxnet中forward的声明如下:

void (const nnvm::NodeAttrs& attrs,
      const OpContext& ctx,
      const std::vector<TBlob>& inputs,
      const std::vector<OpReqType>& req,
      const std::vector<TBlob>& outputs);

下面是forward函数的实现,我们逐行讲解.

template<typename xpu>                                                        // 1
void QuadraticOpForward(const nnvm::NodeAttrs& attrs,                         // 2
                        const OpContext& ctx,                                 // 3
                        const std::vector<TBlob>& inputs,                     // 4
                        const std::vector<OpReqType>& req,                    // 5
                        const std::vector<TBlob>& outputs) {                  // 6
  CHECK_EQ(inputs.size(), 1U);                                                // 7
  CHECK_EQ(outputs.size(), 1U);                                               // 8
  CHECK_EQ(req.size(), 1U);                                                   // 9
  mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();                            // 10
  const TBlob& in_data = inputs[0];                                           // 11
  const TBlob& out_data = outputs[0];                                         // 12
  const QuadraticParam& param = nnvm::get<QuadraticParam>(attrs.parsed);      // 13
  using namespace mxnet_op;                                                   // 14
  MSHADOW_TYPE_SWITCH(out_data.type_flag_, DType, {                           // 15
    MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {                               // 16
      Kernel<quadratic_forward<req_type>, xpu>::Launch(                       // 17
          s, out_data.Size(), out_data.dptr<DType>(), in_data.dptr<DType>(),  // 18
          param.a, param.b, param.c);                                         // 19
    });                                                                       // 20
  });                                                                         // 21
}                                               

• Line 1: xpu 代表设备类型,通过设置cpu和gpu实例化,上述代码可以同时支持CPU和GPU. operator在cc和cu文件中注册时会进行实例化.
• Line 2: attrs包含用户输入的参数a,b和c, 它是整个计算图中一个节点的属性.
• Line 3: ctx包含了称为stream的量,以实现异步执行.stream的作用是保证这段代码在多个GPU上运行次序和单个CPU上一致.
• Line 4: inputs是输入tensor的向量
• Line 5: req是OpReqType的向量,定义结果以何种方式写入输出tensor. req中元素个数和output必须一致.mxnet目前支持三种req类型:null,write和add. null表示跳过计算,write表示用当前operator的结果覆盖对应的输出tensor,add表示把当前的结果和对应输出的tensor相加后保存在输出tensor中. null和add一般在反向传播中使用, 前者用来跳过一些固定参数的梯度计算(比如索引数组),后者用来累计梯度.
• Line 6: output是输出tensor的向量
• Line 7-9:验证每个向量的尺寸
• Line 10: 从ctx中提取stream用来启动kernels
• Line 11-12: 定义输入/输出tensor的引用,方便后续编码. TBlob是可以描述包含任意维度的数据结构,这样不同维度的tensor才可以都放在一个容器中,比如std::vector和std::list. 通过接口get_with_shape可以获得指定维度的TBlob.
• Line 13: 从节点属性中提取用户输入参数a,b和c
• Line 15-21: operator对应的数学公式实现的代码.MSHADOW_TYPE_SWITCH和MXNET_ASSIGN_REQ_SWITCH令代码端支持所有的数据类型和req类型。最内层的宏里，我们启动kernel计算输出tensor，每个线程从输入tensor取一个元素，送入二次函数并把结果按照req类型写入输出tensor。Kernel::Launch只一个通用接口，支持CPU和GPU上的并行计算，保证大部分operator可以用同一份代码运行在CPU和GPU上。kernel函数的定义在下面给出，Map函数就是线程函数。关于kernel struct中关键的两个宏:(1)MSHADOW_XINLINE为CPU和GPU提供统一接口，CPU和GPU编译器都支持这种内联函数。(2)KERNEL_ASSIGN为不同的req类型提供了统一接口。其名字来自于其在并行的计算核心上运行。在CPU上，kernel一般利用OpenMP parallel命令；在GPU上，它就是CUDA的kernel函数。

template<int req>
struct quadratic_forward {
  template<typename DType>
  MSHADOW_XINLINE static void Map(int i, DType* out_data, const DType* in_data,
                                  const float a, const float b, const float c) {
    KERNEL_ASSIGN(out_data[i], req, in_data[i] * (a * in_data[i] + b) + c);
  }
};

Backward Function

Backward函数负责把损失函数相对于最后一层的梯度传播到第一层，这个过程就是反向传播。这里我们不涉及BP的原理和细节，因为网上有很多相关资源。quadratic operator这里要解决的问题是：已知损失函数相对quadratic operator输出的梯度，计算损失函数相对于输入的梯度。因为参数a，b和c不需要调节，所以也不需要计算相对他们的梯度。该问题的数学形式：已知$dL/dy$和$y=a*x^2+b*x+c$,其中$L$表示损失函数，$y$表示quadratic输出tensor，计算目标是$dL/dx$，利用链式法则有

$$dL/dx = dL/dy * dy/dx = dL/dy * (2*a*x+b)$$
上式说明 $dL/dx$依赖于相对输出值的梯度和输入tensor。backward的声明和forward一样，下面是backward的一个实现

template<typename xpu>                                                       // 1
void QuadraticOpBackward(const nnvm::NodeAttrs& attrs,                       // 2
                         const OpContext& ctx,                               // 3
                         const std::vector<TBlob>& inputs,                   // 4
                         const std::vector<OpReqType>& req,                  // 5
                         const std::vector<TBlob>& outputs) {                // 6
  CHECK_EQ(inputs.size(), 2U);                                               // 7
  CHECK_EQ(outputs.size(), 1U);                                              // 8
  CHECK_EQ(req.size(), 1U);                                                  // 9
  mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();                           // 10
  const TBlob& out_grad = inputs[0];                                         // 11
  const TBlob& in_data = inputs[1];                                          // 12
  const TBlob& in_grad = outputs[0];                                         // 13
  const QuadraticParam& param = nnvm::get<QuadraticParam>(attrs.parsed);     // 14
  using namespace mxnet_op;                                                  // 15
  MSHADOW_TYPE_SWITCH(out_grad.type_flag_, DType, {                          // 16
    MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {                              // 17
      Kernel<quadratic_backward<req_type>, xpu>::Launch(                     // 18
          s, in_grad.Size(), in_grad.dptr<DType>(), out_grad.dptr<DType>(),  // 19
          in_data.dptr<DType>(), param.a, param.b);                          // 20
    });                                                                      // 21
  });                                                                        // 22
}                                 

• Line 1-6: backward的声明，和forward一样
• Line 7-9: 检查函数参数的尺寸。因为相对输入的梯度即依赖于相对于输出的梯度也依赖于输入本身，所以input需要包含两个TBlob（译注：如何让调用者知道需要传递两个TBlog进来？查看后面registeration的line 15）
• Line 10: 从ctx中提取stream，后续进行异步执行
• Line 11-13: 设置变量引用，方便后续编码。out_grad是相对于输出的梯度，in_data是输入tensor，in_grad是相对输入的梯度（译注：即待求的梯度）
• Line 14: 获得输入参数 QuadraticParam
• Lines 16-22: 和forward类似，这里是并行计算in_grad的位置。quadratic_backward实现了计算in_grad的公式。

template<int req>
struct quadratic_backward {
  template<typename DType>
  MSHADOW_XINLINE static void Map(int i, DType* in_grad, const DType* out_grad,
                                  const DType* in_data, const float a, const float b) {
    KERNEL_ASSIGN(in_grad[i], req, out_grad[i] * (2 * a * in_data[i] + b));
  }
};

Operator Registration

至此，我们实现了quadratic operator必需的数据结构和函数，下面利用nnvm注册quardtic到前端。注册过程可以理解生成一个operator的实例，存储在operator manager并为这个实例设置属性。
下面的例子来自quadratic_op.cc，负责在CPU上注册这个operator

DMLC_REGISTER_PARAMETER(QuadraticParam);                                           // 1

NNVM_REGISTER_OP(quadratic)                                                        // 2
.describe(R"code(This operators implements the quadratic function:                 // 3
.. math::

    f(x) = ax^2+bx+c

where :math:`x` is an input tensor and all operations
in the function are element-wise.

Example::
  x = [[1, 2], [3, 4]]
  y = quadratic(data=x, a=1, b=2, c=3)
  y = [[6, 11], [18, 27]]

)code" ADD_FILELINE)                                                               // 4
.set_attr_parser(ParamParser<QuadraticParam>)                                      // 5
.set_num_inputs(1)                                                                 // 6
.set_num_outputs(1)                                                                // 7
.set_attr<nnvm::FListInputNames>("FListInputNames",                                // 8
  [](const NodeAttrs& attrs) {                                                     // 9
    return std::vector<std::string>{"data"};                                       // 10
  })                                                                               // 11
.set_attr<nnvm::FInferShape>("FInferShape", QuadraticOpShape)                      // 12
.set_attr<nnvm::FInferType>("FInferType", QuadraticOpType)                         // 13
.set_attr<FCompute>("FCompute", QuadraticOpForward<cpu>)                      // 14
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_quadratic"})  // 15
.set_attr<nnvm::FInplaceOption>("FInplaceOption",                                  // 16
  [](const NodeAttrs& attrs) {                                                     // 17
    return std::vector<std::pair<int, int> >{{0, 0}};                              // 18
  })                                                                               // 19
.add_argument("data", "NDArray-or-Symbol", "Input ndarray")                        // 20
.add_arguments(QuadraticParam::__FIELDS__());                                      // 21

NNVM_REGISTER_OP(_backward_quadratic)                                              // 22
.set_attr_parser(ParamParser<QuadraticParam>)                                      // 23
.set_num_inputs(2)                                                                 // 24
.set_num_outputs(1)                                                                // 25
.set_attr<nnvm::TIsBackward>("TIsBackward", true)                                  // 26
.set_attr<FCompute>("FCompute", QuadraticOpBackward<cpu>);                    // 27

• Line 1: 注册parameter struct
• Line 2: 注册一个名为quadratic的operator。产生一个op实例，存储到operator manager中，并返回刚刚产生的实例的引用。
• Line 3-4: 添加关于operator的说明和用例。documentation engine将提取出这些文档信息在documentation web page上显示
• Line 5: 为operator设置parameter struct解析器。前端利用它解析出参数a，b和c
• Line 6: 设置operator的输入参数个数
• Line 7: 设置operator的输出参数个数
• Line 8-11: 定义函数用来生成输入参数的名称list。这个函数负责在生成operator符号时补充用户没有提供的参数。例如

quad_func=mx.sym.quadratic()

是合法的，因为我们在计算图的operator节点中设置了FListInputName属性。mxnet可以补充上丢失的参数，名称是quadratic0_data，其中quadratic0是operator的名字加上一个序号，而data是FListInputName函数的返回值。用户也可以用如下方式执行quand_func：

quand_exe = quand_func.simple_bind(ctx=mx.cpu(), quandratic0_data=(1,))

• Line 12: 注册shape inference 函数
• Line 13: 注册type inference 函数
• Line 14: 注册forward 函数
• Line 15: 注册backward函数。这里我们使用一个ElemwiseGradUseIn struct。如名字所示，这个注册的函数在计算图增加一个计算梯度的节点，依赖于输入和相对于输出的梯度。mxnet还定义了另外三个类似的struct，ElemwiseGradUseOut，ElemwiseGraduseInOut和ElemwiseGradUseNone。
• Lines 16-19: 这里注册的函数指示了输出tensor可以重用输入tensor的内存，而避免重新分配。operator quandratic只有一个输入和一个输出，所以我们设置了一对0，作为函数的返回值，代表output[0]可以使用input[0]的内存。这里只是给计算图提供一个建议，如果其他节点依赖于输入tensor，则输入tensor的内存不会被覆盖
• Line 20: 定义输入参数的名字为data
• Line 21: 把输入参数a，b和c作为operator的属性
• Line 22: 注册一个名为backward_quadratic的节点，负责operator quadratic的反向计算。名字开头的下划线说明这个operator不会暴露给用户。内部使用反向operator命名习惯是_backward加上对应的forward名字
• Line 23:为_backward_quadratic设置参数解析器
• Line 24:设置输入参数个数
• Line 25:设置输出参数个数
• Line 26:为operator增加TIsBackward属性。shape/type inference利用这个属性确定计算中的节点是forward，还是backward
• Line 27:注册backward函数
  至此我们获得一个operator，可以通过前端调用运行在CPU上。为了运行在GPU上，我们只需要把下面的代码放到quadratic_op.cu中。Note that forward and backward function are registered with attribute key FCompute, rather than FCompute.

NNVM_REGISTER_OP(quadratic)
.set_attr<FCompute>("FCompute", QuadraticOpForward<gpu>);

NNVM_REGISTER_OP(_backward_quadratic)
.set_attr<FCompute>("FCompute", QuadraticOpBackward<gpu>);

Unit Test

我们已经在mxnet后端实现了operator quadratic。如果使用python，当执行
import mxnet as mx时，在线生成两个python函数：一个用于命令式编程，注册为
mxnet.ndarray.quadratic,缩写成mxnet.nd.quadratic;另一个用于符号式编程，注册为mxnet.symbol.quadratic，缩写为mxnet.sym.qudratic。
我们编写如下的test_operator.py文件，在前端测试operator。测试forward很方便，直接利用mx.nd.gradratic即可。但测试backward需要额外的工作量，创建一个quadratic符号，送入辅助函数check_numeric_gradient.辅助函数利用finite difference method，给输入增加小的扰动，计算输出的变化，通过对比来自backward的梯度和finite difference method的梯度值，如果满足用户设定的相对/绝对阈值，则测试通过。

def test_quadratic_function():
    def f(x, a, b, c):
        return a * x**2 + b * x + c

    a = np.random.random_sample()
    b = np.random.random_sample()
    c = np.random.random_sample()
    for ndim in range(1, 6):
        # check forward
        shape = rand_shape_nd(ndim, 5)
        data = rand_ndarray(shape=shape, stype='default')
        data_np = data.asnumpy()
        expected = f(data_np, a, b, c)
        output = mx.nd.quadratic(data, a=a, b=b, c=c)
        assert_almost_equal(output.asnumpy(), expected)

        # check backward using finite difference
        data = mx.sym.Variable('data')
        quad_sym = mx.sym.quadratic(data=data, a=a, b=b, c=c)
        check_numeric_gradient(quad_sym, [data_np])

这里我们用mx.nd.quadratic测试forward函数，check_numeric_gradient用来测试backward函数。mxnet中，另外两个常用的辅助函数有: check_symbolic_forward和check_symblic_backward。利用单元测试中使用，用户需要传入operator符号和期望的结果用来比较。我们推荐对每一个支持backword的operator进行check_numeric_gradient测试，因为check_symbolic_backward可能会因为传入的期望结果不正确而影响测试结果。