0%

MXNet 源码详解--自定义Operator(nnvm方法)

在上一篇中介绍了如何使用 MXNet 的 legacy 接口实现一个Operator, 本文介绍如何使用nnvm接口实现Operator.

本文仍然以实现一个量化训练的Operator为例, 具体的数学过程参考上一篇.
使用nnvm实现Operator和原来的legacy接口要实现的内容比较类似, 只不过nnvm方式所有的接口都是使用函数实现的, 因此, 在Operator中无法记录Operator的中间状态, 所有需要记录的状态都要在外部初始化, 然后以参数的形式传到相应的函数中进行更新.
nnvm 的方法实现Operator同样需要实现3个主要的部分, Parameter 的定义, Operator的计算, 以及Operator的属性.
由于量化训练的原理在上一篇中已经讲过了, 因此, 在这片文章中只关注具体的代码实现.

属性参数定义

属性参数Parameter的定义和之前的完全一致, 这里直接跳过了.

计算的定义

计算的定义接口和legacy的方法也几乎相同, 只不过, 这里不需要继承Operator类了, 而是直接实现具体的计算函数.

Forward

这里为了简单, 同样默认只支持了 float 类型.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
template <typename xpu>
static void QuantizationCompute(const nnvm::NodeAttrs& attrs,
const OpContext& ctx,
const std::vector<TBlob>& inputs,
const std::vector<OpReqType>& req,
const std::vector<TBlob>& outputs) {
using namespace mshadow;
using namespace mshadow::expr;
Stream<xpu>* s = ctx.get_stream<xpu>();
const auto& param = dmlc::get<QuantiParam>(attrs.parsed);
using DType = float;
DType qmin = -127, qmax = 127;
Tensor<xpu, 4, DType> data = inputs[quanti::kData].get<xpu, 4, DType>(s);
Tensor<xpu, 4, DType> out = outputs[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;
}

Backward

在反向中, 需要特别注意的一点是, 对于网络来说, 输出的节点的梯度是放在 inputs 中的, 而输入节点的 grad 是放在 ouputs 中的, 这是因为, 在 nnvm 中, 由于计算图的关系, 在计算梯度的时候, 我们是从前向图的输出节点的grad去计算输入节点的grad, 因此, 计算梯度的计算图恰好和前向图是反的. 在gradient graph 中, 原始数据的输入对应的梯度体现在计算图中是输出, 原始数据的输出对应的梯度体现在计算图中是出入.

1
2
3
4
5
6
7
8
9
10
11
12
13
template <typename xpu>
static void QuantizationGrad(const nnvm::NodeAttrs& attrs,
const OpContext& ctx,
const std::vector<TBlob>& inputs,
const std::vector<OpReqType>& req,
const std::vector<TBlob>& outputs) {
using namespace mshadow;
Stream<xpu>* s = ctx.get_stream<xpu>();
using DType=float;
Tensor<xpu, 4, DType> dgrad = outputs[quanti::kData].get<xpu, 4, DType>(s);
Tensor<xpu, 4, DType> ograd = inputs[quanti::kOut].get<xpu, 4, DType>(s);
Assign(dgrad, req[quanti::kData], ograd);
}

计算属性

同样的, 我们也需要对计算定义一些属性, 类似上一篇中的 InferShape, InferType 等, nnvm 简化了这部分的工作, 直接注册成对应的方法就可以了.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
NNVM_REGISTER_OP(Quanti)
.set_num_inputs(1)
.set_num_outputs(1)
.set_attr_parser(ParamParser<QuantiParam>)
.set_attr<nnvm::FListInputNames>("FListInputNames",
[](const NodeAttrs& attrs) {
return std::vector<std::string>({"data"});
})
.set_attr<mxnet::FInferShape>("FInferShape", ElemwiseShape<-1, 1>)
.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<-1, 1>)
.set_attr<FCompute>("FCompute<cpu>", QuantizationCompute<cpu>)
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_backward_Quanti"})
.add_argument("data", "NDArray-or-Symbol", "")
.add_arguments(QuantiParam::__FIELDS__());
NNVM_REGISTER_OP(_backward_Quanti)
.set_attr_parser(ParamParser<QuantiParam>)
.set_num_inputs(1)
.set_num_outputs(1)
.set_attr<nnvm::TIsBackward>("TIsBackward", true)
.set_attr<FCompute>("FCompute<cpu>", QuantizationGrad<cpu>);

从上面的代码可以看到几个比较重要的地方: InferShape, InferType 和上一篇一样, 也是通过输入的Shape和Type去推导其它输入输出数据的Shape和Type, FCompute用来指定该Operator的具体的计算的实现. 整体上看, 和上一篇基于类的方法要完成的工作是一样的, 只不过换了一种方式.
最后在对应的 .cu 文件中指定在GPU上如何计算之后, 整个Operator就完成了.

1
2
NNVM_REGISTER_OP(Quanti).set_attr<FCompute>("FCompute<gpu>", QuantizationCompute<gpu>);
NNVM_REGISTER_OP(_backward_Quanti).set_attr<FCompute>("FCompute<gpu>", QuantizationGrad<gpu>);

如何选择

最后说一说在实现具体的Operator的时候, 是选择基于上一篇介绍的基于class的方法还是选择本篇介绍的基于nnvm的方法.
理论上, 在写一个Operator的时候, 使用两种方法的任意一种都是可以的, MXNet 需要做到兼容上一篇的class的方法. 两种方法各有优缺点, 但是, 我个人推荐使用nnvm的这种方法.

  1. 由于nnvm兼容之前的代码做的还不够完备, 如果现在使用class方法实现Operator的话可能会出现一些不可预料的问题. 下面会举例说明我遇到的一些兼容方面的问题.
  2. 使用 nnvm 的方法的话Operator写起来更方面一些. 模式也更加统一.

nnvm 兼容旧代码的问题

到目前为止, 我至少碰到了三次因为兼容做的不完备导致而踩到的坑.
一次是 InplaceOption 的坑, 由于没有兼容InplaceOption选项, 网络在设置OpReqType的时候出现了问题, 训练的时候直接崩溃. 解决方法是在 legacy_op_util.cc 中加上对于 InplaceOption 的兼容.
第二次的问题就更隐蔽了, 在使用CacheOp的时候, 由于我需要在运行之前, 根据用户的设置, 修改graph的一些属性, 因此, 在我拿到CacheOp的graph之后, 我会去修改graph中节点的一些attrs并且保存该graph. 但是, 由于兼容的问题, 在我修改了graph之后, 虽然落盘的graph显示是已经修改了attrs, 但是, 实际上在内存中用于训练的graph并没有变更attrs. 最终, 导致, 模型训练完之后再load模型进行inference的时候结果完全不符合预期. 解决方法大概有两个, 一个是把graph保存成json文件之后重新load该文件去build graph, 另外一种方法是在MXNet中fix, 见 https://github.com/shuokay/incubator-mxnet/commit/9865ce9213af956bb200eeb414daff31055d9096
第三次是, 由于nnvm的模式不保存内部状态, 因此导致, 之前用class方法写的带有内部状态的Operator全部都需要做相应的修改, 去掉内部状态.