MXNet new layer(part 1, register: SimpleOp)

Introduction

之前用python写了些新层的实现，后面打算看看C++的实现。把前几天做的整理下，还有比较多的问题待解决。官方的介绍看这里和这里。
note

发现第一部分的内容有些多了，看来要写章回体了…

注册

Simple Op

src/operator/tensor/ 中的文件.cc使用< cpu>，.cu为 < gpu>， .h提供统一的template和具体实现方式,各源文件只是用模板进行注册。有些意思的是.cu（e.g. matrix_op.cu）文件里面在只是简单地注册了GPU的方法，这很nice。但官方的意思是这是种simple Op?
Shape

先来看看关于shape的注册。

    inline bool FlattenShape(const nnvm::NodeAttrs& attrs,
    std::vector<TShape> *in_attrs,std::vector<TShape> *out_attrs) {
    CHECK_EQ(in_attrs->size(),1) << "Input: [data]";
    CHECK_EQ(out_attrs->size(), 1);
    const TShape &dshape = (*in_attrs)[0];
    if (dshape.ndim() == 0) return false;
    out_attrs->clear();
    uint32_t target_dim = 1;
    for (uint32_t i = 1; i < dshape.ndim(); ++i) {
        target_dim *= dshape[i];
    }
    out_attrs->push_back(mshadow::Shape2(dshape[0], target_dim));
    return true;
    }

原型里面有关于TShape的内容,TShape来源于MShadow:

dynamic shape class that can hold shape of arbitrary dimension

# include < tensor_blob.h>

再来关注下倒数第二句，看下Shape2是个什么：

construct a two dimension shape, stride will equal s0

不是很明了，顺带往上看了Shape1:

MSHADOW_XINLINE Shape < 1 > mshadow::Shape1 ( index_t s0 )
construct a one dimension shape, stride will equal s0
Parameters
s0 size of dimension 0

于是明白了，dshape[0] 大致对应了batch_channel,这一点可以从for循环的起始数得到印证。
ElemwiseShape,ElemwiseType

关于这两个函数，官网上没给太多的解释。从程序上来看（matrix_op.cc有很多例子），应该是输入输出保持一致的意思，比如：

NNVM_REGISTER_OP(Flatten)
...
.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>)
...

NNVM_REGISTER_OP(dot)
...
.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<2, 1>)
...

再配合下官网解释：

Use ElemwiseShape< n_in, n_out> for simple operators with uniform shapes.

左边是输入参数的个数，右边是输出参数。

再看下另外一个op：

NNVM_REGISTER_OP(flip)
.MXNET_DESCRIBE("Flip the input tensor along axis and return a new one.")
...
.set_attr<nnvm::FInferShape>("FInferShape", ElemwiseShape<1, 1>)
.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>)
...

也就是说ElemwiseShape和ElemwiseType表示输入输出一致，这也就是官网上说的uniform的意思。但这会有个问题，n_in和n_out的含义？是不是可以支持指定那些是一致的？看来还是要找到源程序才行。

// src/operator/elemwise_op_common.h
...
template<int n_in, int n_out>
inline bool ElemwiseShape(const nnvm::NodeAttrs& attrs,
                          std::vector<TShape> *in_attrs,
                          std::vector<TShape> *out_attrs) {
  CHECK_EQ(in_attrs->size(), n_in) << " in operator " << attrs.name;
  CHECK_EQ(out_attrs->size(), n_out);
  return ElemwiseAttr<TShape, shape_is_none, true>(                                                                            
    attrs, in_attrs, out_attrs);
}
...

template<int n_in, int n_out>                                                                                                  
inline bool ElemwiseType(const nnvm::NodeAttrs& attrs,
                         std::vector<int> *in_attrs,
                         std::vector<int> *out_attrs) {
  CHECK_EQ(in_attrs->size(), n_in) << " in operator " << attrs.name;
  CHECK_EQ(out_attrs->size(), n_out);
  return ElemwiseAttr<int, type_is_none, true>(
    attrs, in_attrs, out_attrs);
}
...

检查了一致性就扔掉了，差不多只是为了提醒程序员。此处注意到TShape参与Shape的工作，用int表示数据类型。

FGradient

然后是FGradient:

NNVM_REGISTER_OP(Flatten)
...
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{ "_backward_copy" })
...

追踪一下：

//elemwise_op_common.h
struct ElemwiseGradUseNone {                                                                                                   
  const char *op_name;
  std::vector<nnvm::NodeEntry> operator()(const nnvm::NodePtr& n,
                                          const std::vector<nnvm::NodeEntry>& ograds) {
    return MakeGradNode(op_name, n, ograds, n->attrs.dict);
  }
};

有意思的是，在这个struct里面提供了一个运算符，后面找时间可以试试。
初看起来后面要为其分配空间，但这样做似乎有些不明智。此处的大环境是注册而已，不会有实际操作。
来看看官网的介绍:

Use utility functions ElemwiseGradUseIn{op_name}, ElemwiseGradUseOut{op_name}, ElemwiseGradUseNone{op_name} for ops that need corresponding forward op’s input, output or nothing to calculating gradient.

不是很具体，再看下一段做个参考:

For more complicated pattern, use MakeGradNode(op_name, n, heads, dict) to create gradient entries, where heads are input entries to the backward op, composed from ograds and n->inputs.

再来些程序:

NNVM_REGISTER_OP(Flatten)
...
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{ "_backward_copy" })
...

，结合elemwise_op_common.h中struct ElemwiseGradUseNone的结构，很容易认为* “_backward_copy”* 只是一个字符串。实则另有玄机,比如换一个operator来看：

NNVM_REGISTER_OP(dot)
...
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_dot"})
...

NNVM_REGISTER_OP(_backward_dot)
...
.set_attr<FCompute>("FCompute<cpu>", DotBackward_<cpu>)
...

note

1. 此处可以看出字符串是与注册的op 一致的，所以并不是简单的一个字符串而已，（前面有提到注册时不是用的字符串类型）;

2. 此处用的是ElemwiseGradUseIn()
   也是可以和前面来个对比了：
   结合前面提及的官网解释，可以猜测出这三种的用法：dot的后向操作需要用到前向操作的输入，而flatten不需要。这还要从程序中找些证据：

   // src/operator/elemwise_op_common.h
   struct ElemwiseGradUseIn 
   {
   const char *op_name;
   std::vector<nnvm::NodeEntry> operator()(const nnvm::NodePtr& n,
                                         const std::vector<nnvm::NodeEntry>& ograds) 
       {
   std::vector<nnvm::NodeEntry> heads(ograds.begin(), ograds.end());
   for (auto& h : n->inputs){
     heads.push_back(h);
           }   
   return MakeGradNode(op_name, n, heads, n->attrs.dict);
       }
   };
   
   struct ElemwiseGradUseOut {
   const char *op_name;
   std::vector<nnvm::NodeEntry> operator()(const nnvm::NodePtr& n,
                                         const std::vector<nnvm::NodeEntry>& ograds) {
   std::vector<nnvm::NodeEntry> heads(ograds.begin(), ograds.end());
   index_t n_out = n->num_outputs();
   for (index_t i = 0; i < n_out; ++i) {
     heads.emplace_back(nnvm::NodeEntry{n, i, 0});
           }   
   return MakeGradNode(op_name, n, heads, n->attrs.dict);
       }
   };
   
   struct ElemwiseGradUseNone {
   const char *op_name;
   std::vector<nnvm::NodeEntry> operator()(const nnvm::NodePtr& n, const std::vector<nnvm::NodeEntry>& ograds) {
   return MakeGradNode(op_name, n, ograds, n->attrs.dict);
       }
   };
   

   此处有些意思的是，UseIn和UseOut使用的添加方式不一样，一个是将已有实体添加进容器，另一个是新建一个，这可能和其运行机理有关，但接口应该是一致的，现阶段还是先省省吧…动态运行放后面去。
   于是，这三中方法是注册了后向方法的参数信息。具体来说，就是定义input和output数据 （对forward操作而言）在backward方法中的分配，而梯度信息ograds是被默认包括的。

   验证一下：

   template<typename xpu> 
   void DotBackward_(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::expr;
   ...
   if (inputs[1].ndim() == 2 && inputs[2].ndim() == 2) { 
   mshadow::Tensor<xpu, 2, real_t> mout_grad = inputs[0].get<xpu, 2, real_t>(s);                                              
   mshadow::Tensor<xpu, 2, real_t> mlhs_data = inputs[1].get<xpu, 2, real_t>(s);
   mshadow::Tensor<xpu, 2, real_t> mrhs_data = inputs[2].get<xpu, 2, real_t>(s);
   mshadow::Tensor<xpu, 2, real_t> mlhs_grad = outputs[0].get<xpu, 2, real_t>(s);
   mshadow::Tensor<xpu, 2, real_t> mrhs_grad = outputs[1].get<xpu, 2, real_t>(s);
   ...
   

   另外提一下forward和backward复用的例子:

    NNVM_REGISTER_OP(flip)
    ...
    .set_attr<FCompute>("FCompute<cpu>", Flip<cpu>)
    .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"flip"})
    ...

有趣吧 :)