Operators in MXNet

MXNet 中的 operator 是一个既包含了实际的计算也包含了一些附加信息的 class, 这些附加信息可以帮助我们的系统来实现原地更新和自动微分等优化. 在在继续这篇文档之前, 我们强烈的建议您首先要搞明白 mshadow 库, 因为所有 operator 的计算都是基于系统提供的数据结构 mshadow::TBob, 而该数据结构类似于张量(Tensor). MXNet 的 operator 的接口致力于灵活性, 现在提供的灵活性包括以下几点:

• 通过特定的原地更新数据来节省内存.
• 在 python 端掩盖一些内部的参数, 是代码整洁.
• 定义输入 tensor 和输出 tensor 的关系,这样系统可以帮助我们进行类型检查.
• 从系统中申请额外的空间来进行计算 (e.g. calling cudnn routines).

Operator Interface

operator 核心的接口是 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;

• OpContext 的数据结构是下面的代码:

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

  , 你可以知道 operator 是在进行 train 还是 test (is_train); operator 应该在哪个device ( run_ctx ) 上运行以及通过下面的参数来确定是否要申请额外的资源.

• in_data 和 out_data 分别代表输入 tensor 和输出 tensor. 所有的 tensor 需要的空间都是系统进行申请和管理.

• req 表示计算的结构是如何写入到 out_data 中的. 换句话说, req.size() == out_data.size() 和 req[i] 与如何写入 out_data[i] 是相关的. OpReqType 是下面这样定义的:

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

, 一般情况下, 所有的 out_data 的类型应该是 kWriteTo, 表示out_data 代表的 tensor 提供的是可以直接写入的 原始的 内存块 . 在有些情况下, 比如说在计算 表示 gradient 的 tensor 的时候, 我们最好是将梯度累加起来, 而不是直接覆盖掉原来的结果, 这样我们就不需要每次计算的时候申请额外需要的内存空间. 在这种情况下, req 的类型应该是 kAddTo, 表示应该调用 += 操作.

• aux_states 表示的是为了方便计算而需要的附加的 tensor, 现在是没有用到的.

除了 Forward operator, 用户有时候也需要实现 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);

Backward 的接口遵循 Forward 一样的设计原则, 除了 out_grad, in_data 和 out_data 是 operator 计算 in_grad 必须的以外, 其他的输入和 Forward 是一样的. 命名策略和 torch 的约定很类似, 可以用下面的图示来总结:

[input/output semantics figure]

有些 operator 并不是都需要的out_grad,in_data和out_data 参数, 这个需求可以通过OperatorProperty中的接口DeclareBackwardDependency来实现.

Operator Property

有这么一种可能, convolution 有好几种不同的实现, 用户可能想要在这些算法中选择能够获得最高性能的算法. 为了实现这个目的, 我们将 operator 的 sematic 接口从具体的实现 (Operator 类) 中分离出来, 独立为OperatorProperty 类. OperatorProperty的接口包括以下内容:

• InferShape:

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

这个接口有两个目的: (1) 向系统提供每个输入和输出 Tensor 的大小, 这样系统可以在进行 Forward 和 Backward 之前提前申请好相应的内存; (2) 进行类型检查, 在运行前确保没有明显的错误. in_shape 中的 shape 是有系统自动设置 (依据是依赖的上个 Operator 的 out_shape ). 这个接口会返回false如果系统认为提供的信息不足以完成shape的推断, 或者在shape 不一致的时候抛出异常.

• Request Resources: 有些操作需要额外的内存作为工作空间来进行计算, 比如说cudnnConvolutionForward. 这种情况下, 系统最好可以对这部分内存进行管理, 这样系统可以做一些优化, 比如说内存的重复利用. MXNet 定义了两个接口来达到目的:

  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 数据结构 (在 resource.h 中) 现在只包含了一个 type flag:

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

如果 ForwardResource 和 BackwardResource 返回的数组是非空的, 那么系统会通过Operator 的Foward 和Backward 接口中的 ctx 参数来提供相应的资源. 简单的举个例子, 如果要获取这些资源, 可以按照下面的写法来做:

auto tmp_space_res = ctx.requested[kTempSpace].get_space(some_shape, some_stream);
auto rand_res = ctx.requested[kRandom].get_random(some_stream);

具体的例子可以参考 src/operator/cudnn_convolution-inl.h.

• Backward dependency: 让我们来看两个不同的 operator signature ( 为了演示的目的，我们命名了所有的变量):

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

我们注意到, 在 FullyConnectedForward 中使用的 out_data变量, 在 FullyConnectedBackward 没有被用到. 与此同时, PoolingBackward 使用了 PoolingForward 中所有的变量. 因此, 这个out_data tensor 在做backward 不需要的时候, 需要释放. 这里有个挑战就是如何进行垃圾回收 (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 元素一个区分不同的 arrays 的 id. 让我们来看看这个接口如何定义FullyConnected和Poolling的不同的依赖关系的:

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]};
}

• Inplace Option: 为了进一步的节省内存的申请开销, 我们倾向于是用原地更新(inplace update). 这个主要用在 element-wise 操作上, 因为这种情况下输入 tensor 和输出 tensor 的 shape 是一致的. 针对这种情况下, 我们提供了下面的接口:

  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] tensors 应该在Forward 的计算过程中使用同样的内存空间, 同样地, out_grad[0] 和 in_grad[0] 分享同样的内存空间在 Backward计算过程中.

ATTENTION: Even with the above specification, it is not guaranteed that input and output tensors will share the same space. In fact, this is only a hint for the system for the final decision. However, in either case, such decision is completely transparent to user, so the actual Forward and Backward implementation does not need to consider that.

• Expose Operator to Python: 由于 c++ 编程语言的限制, 我们需要用户实现下面的接口:

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

Create Operator from Operator Property

我们在上面的内容中提到过 OperatorProperty 包括所有的一个操作的 semanticalattributes. 它也包括需要创建一个Operator 指针指向真正的计算操作.

Create Operator

实现 OperatorProperty中下面的接口:

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

For example:

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

Parameterize Operator

当我们实现一个卷积operator 的时候, 我们需要知道kernal 大小, stride 大小, padding 大小等等信息. 这些需要作为参数传递给 operator, Forward 和Backward计算过程需要这些参数. 为了传递参数, 用户需要定义 ConvolutionParam 数据结构:

#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里, 然后在类初始化的时候, 将这个数据结构传递给operator 类中.

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_;
};

Register Operator to MXNet

使用下面的宏定义来将 parameter 结构和 OperatorProperty 类注册到 MXNet 的系统中:

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

这个宏定义的第一个参数是 name, 第二个参数是 Property 类的名字.

All in a list

Finally! We almost covered the interface we needed to define a new operator. Let’s do a recap in a list:

• Use Operator interface to write your actual computation logic (Forward and Backward).
• Use OperatorProperty interface to:
  • Pass parameter to operator class (may use Init interface).
  • Create operator using CreateOperator interface.
  • Correctly implement the operator description interface such as the names of arguments, etc.
  • Correctly implement the InferShape interface to set the output tensor shape.
  • [Optional] If additional resources are needed, check ForwardResource and BackwardResource.
  • [Optional] If Backward does not need all the input and output of Forward, check DeclareBackwardDependency.
  • [Optional] If inplace update is supported, check ForwardInplaceOption and BackwardInplaceOption.
• Register the OperatorProperty class and the parameter class.

祝您 MXNet 之旅愉快.