本文是MXNet源码解读系列的第二篇, 本篇介绍如何使用C++实现新的Operator.

随着MXNet的演进, 在MXNet中实现新的Operator有两种方式, 一种是早期的版本, 通过定义计算类和相应的属性类的方式, 另外一种是现在MXNet中在使用的通过nnvm的方式. MXNet仍然对早期通过类定义的Operator的方式提供良好的兼容. 本文将详细介绍如何通过类的方式实现新的Operator. 具体来说本文通过实现一个简单的量化训练Operator来详细解释如何在MXNet中实现一个新的Operator.

量化训练

关于量化训练的详细内容可以参考其它的文章, 简单来说, 本文将实现的Operator为:
$$
y = \text{clip}(\text{round}(s \cdot x), -127, 127)
$$
在MXNet中实现一个新的Operator需要实现3个主要的部分, Operator的属性参数, Operator的计算, Operator 的属性信息. 接下来逐个介绍.

属性参数定义

Operator 的属性参数定义了Operator的一般属性, 例如, 在Convolution中的, kernel_size, num_filter, strides, padding 等等这些信息, 都属于属性参数. 在本文的量化训练的例子中, scale 是属性参数. 属性的参数定义通过继承 dmlc::Parameter来实现.

struct QuantiParam : dmlc::Parameter<QuantiParam> {
  float scale;
  DMLC_DECLARE_PARAMETER(QuantiParam) {
    DMLC_DECLARE_FIELD(scale).describe("the scale of quantization");
  }
};

关于Parameter内部的实现, 可以参考源码或者关注后续的文章, 当前可以先照猫画虎实现对应的Parameter. 其中, 要注意的一点是, 在Parameter中定义的所有的参数, 都必须支持 >> 这个操作符, 这是因为, 在MXNet的传递中使用的是 std::pair<std::string, std::string>, 在Parameter内部解析的时候, 需要把的std::string 类型转换成对应的参数类型, 底层是通过 >> 这个操作符实现的.

计算的定义

对于一个Operator来说, 最重要的就是定义该Operator的计算是什么, 在MXNet中, 我们定义的Operator通常都是要用来训练神经网络的, 因此, 我们要同时定义 Forward 和 Backward, 如果明确不需要Backward的时候,也可以不实现Backward. 例如, 在部署量化模型的预测库中, 因为只需要跑前向计算, 因此, 我们就可以只实现Forward而不用实现Backward. 这里, 因为我定义的是量化训练的Operator, 因此, 我会同时实现Forward和Backward.

Forward

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) {
using namespace mshadow;
using namespace mshadow::expr;
Stream<xpu>* s = ctx.get_stream<xpu>();
DType qmin = -127, qmax = 127;
Tensor<xpu, 4, DType> data = in_data[quanti::kData].get<xpu, 4, DType>(s);
Tensor<xpu, 4, DType> out = out_data[quanti::kOut].get<xpu, 4, DType>(s);
DType scale = param_.scale;
out = F<mshadow_op::minimum>(F<mshadow_op::maximum>(F<mshadow_op::round>(scale * data), qmin),
                                qmax) / scale;
}

首先我们关注Forward的接口, ctx 是用来定义Operator是在什么设备上跑的, 例如 CPU还是GPU, in_data 是该计算所需要输入的数据, 例如, 在Convolution中, in_data 就是计算卷积 data, weight 和 bias, 在本文的量化Operator中是输入数据data, 也就是上述公式中的 x, req 是用来提示memory的复用关系, 它定义了Operator的计算结果是如何放到out_data中去的, 这个参数基本上不需要用户干预, 框架会自己维护该参数. out_data 是计算的输出. aux_states 是辅助参数, 例如, 在BatchNorm中的running_mean, running_var 等都是 aux_states, 简单来说, aux_states 是模型的参数, 但是, 他和 in_data 中的参数的不同点是, aux_states 不需要计算梯度.
第10-11行: 把 TBlob 类型的输入数据转换成 mshadow::Tensor 类型, 其目的是方便使用 mshadow 进行计算. Tensor 并没有存储具体的数据, 只是保存了指向数据的指针, 因此, 这个转换和具体的计算相比, 其代价可以忽略不计. mshadow 是一个实现了lazy compute的Tensor计算库, 在MXNet中, 熟练掌握shadow能给实现Operator带来巨大的便利.
第13-14行: 量化训练的实现, F 是 mshadow 中定义的一个接口, 在这里我们完全使用了mshadow完成了具体的计算, 并且, 通过mshadow实现的代码是可以同时在CPU和GPU上面运行的, 因此, 我们不需要单独再实现GPU计算的代码了.

Backward

因为我们要实现量化训练, 因此, 这里我们还需要实现对应的 Backward 的计算. 我们这里实现一个最简单的版本–其反向的梯度永远是 1.
$$
\text{grad} = 1
$$

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) {
using namespace mshadow;
Stream<xpu>* s = ctx.get_stream<xpu>();
Tensor<xpu, 4, DType> dgrad = in_grad[quanti::kData].get<xpu, 4, DType>(s);
Tensor<xpu, 4, DType> ograd = out_grad[quanti::kOut].get<xpu, 4, DType>(s);
Assign(dgrad, req[quanti::kData], ograd);
}

Forward的计算中我们没有展示req是如何使用的, 在这里, 我特意使用了 req 来表明计算出来的梯度要如何放到 dgrad 中. 如果深究的话, OpReqType 有以下几个选项: kNullOp(什么也不做), kWriteTo(把梯度写到 dgrad 中), kWriteInplace(原地写, 这个要配合Operator属性定义中的InplaceOption使用), kAddTo(和原来的数据相加)

属性定义

在QuantiParam中我们已经定义了Operator的一些属性, 在QuantiParam中定义的参数, 最终都会暴露给用户, 用户可以根据具体的需求给参数设置不同的值. 这里定义的属性, 和Operator的计算行为没有关系, 不管这里如何定义, 都不应该影响Operator的计算行为(如果影响了, 那么说明代码里有bug). 这里的属性, 有些是需要在运行时才知道具体行为的, OperatorProperty 中包含了Operator的所有的信息. OperatorProperty 的接口很多, 但是, 大部分时间我们不需要全部实现. 这里, 为了简单, 我们只实现几个必要的接口. 其它的接口有的是为了进行memory优化的, 有的是用来定义Operator的数据类型的, 可以在需要的时候做具体的实现, 尤其是InferType接口, 在实现量化预测库的时候, 因为所有的数据都是使用的不同位宽的整型数据, 在实现的时候需要在InferType中做具体的推断.

class QuantiProp : public OperatorProperty {
 public:
  void Init(const std::vector<std::pair<std::string, std::string>>& kwargs) override {
    param_.Init(kwargs);
  }
  std::map<std::string, std::string> GetParams() const override { return param_.__DICT__(); }
  bool InferShape(std::vector<TShape>* in_shape,
                  std::vector<TShape>* out_shape,
                  std::vector<TShape>* aux_shape) const override {
    SHAPE_ASSIGN_CHECK(*out_shape, quanti::kOut, (*in_shape)[quanti::kData]);
    return true;
  }

  OperatorProperty* Copy() const override {
    auto ptr = new QuantiProp();
    ptr->param_ = param_;
    return ptr;
  }

  std::string TypeString() const override { return "Quanti"; }

  Operator* CreateOperator(Context ctx) const override {
    LOG(FATAL) << "Not Implemented.";
    return NULL;
  }

  Operator* CreateOperatorEx(Context ctx,
                             std::vector<TShape>* in_shape,
                             std::vector<int>* in_type) const override;

 private:
  QuantiParam param_;
};

在这里最重要的就是 InerShape 这个接口, 该接口给定第一个输入的shape, 然后根据QuantiParam参数信息推导出其它输入的shape和输出的shape. 本文的Operator非常简单, 只有一个输入和一个输出, 并且输入和输出的shape完全一样, 因此, 该接口实现起来非常简单. 在其它的一些Operator中, 例如, Convolution中, 输入有 data, weight, bias(也有可能没有), 并且他们的shape都不一样, 这时候就可以在这里通过data的shape和Convolution的Parameter参数推断出weight, bias, output的shape.

注册Operator

以上我们实现了具体的Operator, 为何运行起来, 我们还需要把Operator注册到框架中.

Operator* QuantiProp::CreateOperatorEx(Context ctx,
                                       std::vector<TShape>* in_shape,
                                       std::vector<int>* in_type) const {
  std::vector<TShape> out_shape, aux_shape;
  std::vector<int> out_type, aux_type;
  CHECK(InferType(in_type, &out_type, &aux_type));
  CHECK(InferShape(in_shape, &out_shape, &aux_shape));
  DO_BIND_DISPATCH(CreateOp, param_, (*in_type)[0]);
}

这个接口实现在大部分Operator中也不需要修改. CreateOp 要根据CPU或者GPU做相应的特化. 例如, 在CPU上的定义为:

template <>
Operator* CreateOp<cpu>(QuantiParam param, int dtype) {
  Operator* op = new QuantiOp<cpu, float>(param);
  return op;
}

另外注意, 这里直接把QuantiOp实例化为了 float 类型, 如果要支持其他类型, 比如 float16, 可以根据 dtype 来做具体的实例化.
最后就是提供给用户的接口:

DMLC_REGISTER_PARAMETER(QuantiParam);
MXNET_REGISTER_OP_PROPERTY(Quanti, QuantiProp)
    .add_argument("data", "Symbol", "Input data to the Quanti")
    .add_arguments(QuantiParam::__FIELDS__())
    .describe("Quanti");

其中, 第 3 行是Operator的计算数据, 第4行是Operator的属性参数.

支持GPU上计算

为了支持GPU计算, 我们还需要创建GPU上对应的Operator, 和上面cpu上的CreateOp类似, 把CreateOp特化到gpu上就好了.

Operator* CreateOp<gpu>(QuantiParam param, int dtype) {
  Operator* op = new QuantiOp<gpu, float>(param);
  return op;
}

总结

本文介绍了如何动手实现MXNet的Operator, 为了重点关注实现一个Operator中的主要工作内容, 本文没有介绍关于memory优化的内容(在很多时候还是比较重要的)以及 OperatorProperty 中其它接口的定义和实现. 本文的主要目的是用最简洁的例子和代码教大家实现一个能运行的Operator, 至于各种优化, 各种特殊场景下的接口实现, 相信大家遇到具体需求的时候会逐渐学会.
本文所有的代码实现在这里: https://github.com/shuokay/incubator-mxnet/tree/legacy-operator/plugin/quantization

Memo

MXNet 源码详解--自定义Operator