MXNet的运算符-Part1

一个运算符Operator是一个类，同时包含计算逻辑和辅助信息(帮助执行性能优化，比如及时更新和自动求导)

为此，建议参考mshadow库,类似tensor结构的mshadown:TBlob

MXNet运算符接口的用处：

通过指定及时更新来减少内存分配

从Python来隐藏内部参数

定义输入和输出tensors的关系，允许系统执行模型检查

从系统中获得附加的临时空间执行计算；比如调用cudnn routines

运算符接口(Operator Inteface)

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) = 0;

上下文

 struct OpContext {
             int is_train;
             RunContext run_ctx;
             std::vector<Resource> requested;
           }

定义：运算符所在的阶段(Train,Test)

         运算符在什么设备上执行，以及所需的资源。

req指示：计算结果如何写入out_data

req.size() == out_data.size() 

并且req[i]对应到out_data[i]

and req[i] correspond to the write type of out_data[i]

运算请求类型OpReqType 

enum OpReqType {
             kNullOp,
             kWriteTo,
             kWriteInplace,
             kAddTo
           };

kWriteTo和kAddTo之间的差异化

通常情况下，所有的out_data都是kWriteTo

但是针对计算梯度下降的Tensor,计算量非常巨大，通常采用递增的方式处理

aux_states用来设计支持辅助的Tensor的，目前用处不大。

Backward接口

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);

该接口类似Forward，除了out_grad,in_grad

注意:in_grad是结果

命名策略类似Torch的规约 [input/output semantics figure]

针对部分运算符，不需要out_grad,in_data,out_data；可以用DeclareBackwardDependency来指定依赖关系。

Operator Property[运算符的属性]

卷积有许多实现，你或者想在这些卷积之间进行网络切换，并获取最好的性能；

为此，分离运算符的语义接口和实现接口；Opeartor 和 OpeartorProperty

关注OperatorProperty的核心

InferShape 推导模型

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

该接口的主要用途

1.告诉每个Input和Output的Tensor的大小，在执行Forward和Backward之前分配好空间。

2.执行大小检查，in_shape的由系统来设置，根据之前的out_shape来完成；如果没有足够多的信息来推导的话，会返回false

如果模型不一致，则抛出异常。

Request Resources

类似cudnnConvolutionForward运算符需要专门的计算空间；如果一个系统能够管理，就可以执行优化，进行重复利用。

virtual std::vector<ResourceRequest> ForwardResource(
               const std::vector<TShape> &in_shape) const;
virtual std::vector<ResourceRequest> BackwardResource(
               const std::vector<TShape> &in_shape) const;

注意结构体 ResourceRequest 的定义

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

如果前向和后向资源返回非空数组；系统会通过接口中的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

关注下两类不同类型的操作符定义

void FullyConnectedForward(TBlob weight, TBlob in_data, TBlob out_data);
void FullyConnectedBackward(TBlob weight, TBlob in_data, TBlob out_grad, TBlob in_grad);

void PoolingForward(TBlob in_data, TBlob out_data);
void PoolingBackward(TBlob in_data, TBlob out_data, TBlob out_grad, TBlob in_grad);

注意FullyConnectedFoward的out_data没有在FullyConnectedBackward中使用

而PoolingBackward会全部使用PoolingForward的参数

结果：FullyConnectedForward的out_data在一次消费后，就可以进行释放了；可以通过GC来完成

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

参数Vector中的int元素，是一个区分不同数组的ID

关注这些接口如何指定FullyConnect和Pooling之间的不同依赖

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

std::vector<int> PoolingProperty::DeclareBackwardDependency(
              const std::vector<int> &out_grad,
              const std::vector<int> &in_data,
              const std::vector<int> &out_data) const {
            return {out_grad[0], in_data[0], out_data[0]};
          }

In Place Option

为了保持内存分配的成本，可以通过In Place Update.

当input和output向量具有相同的模型时，对于element-wise的运算符是恰当的；

virtual std::vector<std::pair<int, void*>>  ElewiseOpProperty::ForwardInplaceOption(
               const std::vector<int> &in_data,
               const std::vector<void*> &out_data) const {
             return { {in_data[0], out_data[0]} };
           }
virtual std::vector<std::pair<int, void*>> ElewiseOpProperty::BackwardInplaceOption(
               const std::vector<int> &out_grad,
               const std::vector<int> &in_data,
               const std::vector<int> &out_data,
               const std::vector<void*> &in_grad) const {
             return { {out_grad[0], in_grad[0]} }
           }

针对系统来说，in_data[0]和out_data[0]在执行Forward时可以共享相同的内存；同理，在Backward时，out_grad[0]和in_grad[0]同样允许。

核心关键：针对依赖的执行，在系统层面是否能够共享内存，针对用户态需要做到透明化。

向Python暴露运算符

// initial the property class from a list of key-value string pairs
virtual void Init(const vector<pair<string, string>> &kwargs) = 0;
// return the parameters in a key-value string map
virtual map<string, string> GetParams() const = 0;
// return the name of arguments (for generating signature in python)
virtual vector<string> ListArguments() const;
// return the name of output values
virtual vector<string> ListOutputs() const;
// return the name of auxiliary states
virtual vector<string> ListAuxiliaryStates() const;
// return the number of output values
virtual int NumOutputs() const;
// return the number of visible outputs
virtual int NumVisibleOutputs() const;

从操作属性中创建一个操作符

OperatorProperty包括一个Operation的所有语法属性，可以针对执行任务时，创建Operator指针。

Create Operator

通过实现OperatorProperty的接口

virtual Operator* CreateOperator(Context ctx) const = 0;

实例

class ConvolutionOp {
     public:
      void Forward( ... ) { ... }
      void Backward( ... ) { ... }
    };
    class ConvolutionOpProperty : public OperatorProperty {
     public:
      Operator* CreateOperator(Context ctx) const {
        return new ConvolutionOp;
      }
    };

参数化运算符

在实现一个卷积运算符时，需要知道kernel size, stride size, padding size等配置参数；

并且在forward和backward之前把这些参数传给Opeartor;可以定义一个卷积参数

#include <dmlc/parameter.h>
    struct ConvolutionParam : public dmlc::Parameter<ConvolutionParam> {
      TShape kernel, stride, pad;
      uint32_t num_filter, num_group, workspace;
      bool no_bias;
    };

把该参数设置在ConvolutionOpProperty中，并且在构造器中传入

class ConvolutionOp {
     public:
      ConvolutionOp(ConvolutionParam p): param_(p) {}
      void Forward( ... ) { ... }
      void Backward( ... ) { ... }
     private:
      ConvolutionParam param_;
    };
    class ConvolutionOpProperty : public OperatorProperty {
     public:
      void Init(const vector<pair<string, string>& kwargs) {
        // initialize param_ using kwargs
      }
      Operator* CreateOperator(Context ctx) const {
        return new ConvolutionOp(param_);
      }
     private:
      ConvolutionParam param_;
    };

把OperatorProperty Class 和Parameter类注册到MXNet

可以采用如下宏来完成

DMLC_REGISTER_PARAMETER(ConvolutionParam);
MXNET_REGISTER_OP_PROPERTY(Convolution, ConvolutionOpProperty);

接口汇总

1.利用Opeartor接口来写forward,backward的计算逻辑

2.使用OpeartorProperty可以满足如下场景的使用

    2.1利用Init接口来传递参数

    2.2用CreateOpeartor接口来创建操作符

3.注册OpeartorProperty类和参数类