神经网络层与操作注册机制调研

MXNet (NNVM)

nnvm是从mxnet的实现中剥离出来一个模块，该模块完成了从symbol描述的网络到graph描述的符号计算图的生成和优化工作，而这样的模块化剥离仿效了unix的哲学，使得mxnet能够在不同的设备应用和场景中自主裁剪各功能模块。

1. 定义

1.1 Node

计算图中的一个Node, 包含以下信息：

• inputs: node的输入（NodePtr）
• attrs: node的一些属性
• Op: 每个node有一个operator

1.2 Operator

一个Op就是一个操作，一个Node对应一个Op。

1.3 Pass

对包含了各种属性映射信息的计算图执行转换，转换使得该计算图拥有更多地属性或者变为另一个计算图。 nnvm中实现的pass包括形状和类型推断、内存计划等。

1.4 Graph

计算图的结构，中间表示方式，不可执行。

• 计算图的所有叶子节点
• 存储每个节点各种属性的map: 一个从字符串到任意类型的属性映射map< string, shared_ptr< any > >，这个属性映射包含了每一个tensor的shape、type以及内存分配计划。any可以是：
  • using ShapeVector = std::vector<TShape>;
  • using DeviceVector = std::vector<int>;
  • using DTypeVector = std::vector<int>;
  • using StorageVector = std::vector<int>;

1.5 IndexGraph

IndexGraph是以一维数组存储图模型的，它根据symbol的拓补排序排列元素，TopoSort函数对输入个数为0的节点进行后向深度优先遍历得到拓补排序。 通过遍历Graph生成IndexGraph。

2. 特性

• 简单的注册机制
• 开发者可以灵活注册attributes of operator
• 对computation graph进行自定义的优化

2.1 简单的注册机制

nnvm采用NNVM_REGISTER_OP去注册一个operator，并可以对不同的op用set_attr注册不同的属性，使得不同operator的实现不必采用同一个operator接口，完成了去中心化的设计目标，使得不同框架下的operator实现都可以采用nnvm做计算图优化。 注册代码如下：

 NNVM_REGISTER_OP(add)
.describe("add two data together")
.set_num_inputs(2);

NNVM_REGISTER_OP(conv2d)
.describe("take 2d convolution of input")
.set_num_inputs(2);

NNVM_REGISTER_OP(assign)
.describe("assign second input argument to the first one")
.set_num_inputs(2);

在python中使用已注册的operator:

import nnvm.symbol as nn

# symbolic variable
x = nn.Variable('x')
y = nn.Variable('y')
w = nn.Variable('w')

z = nn.conv2d(nn.add(x, y), w, filter_size=(2,2), name='conv1')

2.2 注册自定义attributes

接口非集中管理，可以方便的进行扩展。 用户可以灵活的为operator注册attribute, 比如：我们为addoperator注册一个function来检查是否可以calculate inplace：

using FInplaceOption = std::function<
  std::vector<std::pair<int, int> > (const NodeAttrs& attrs)>;

// attributes can be registered from multiple places.
NNVM_REGISTER_OP(add)
.set_num_inputs(1);

// register to tell first input can be calculate inplace with first output
NNVM_REGISTER_OP(add)
.set_attr<FInplaceOption>("FInplaceOption", [](const NodeAttrs& attrs) {
  return std::vector<std::pair<int, int> >{{0, 0}};
 });

2.3 computation graph优化

通过一个或多个pass对computation graph进行优化。使用pass的方法如下：

Graph ApplyPasses(Graph src, const std::vector<std::string>& passes);

一个Graph经过pass处理（加工？）成新的Graph，新的Graph会包含更多attributes信息。比如：经过plan_memory pass的处理，Graph会增加内存分配计划相关的attribute.

3. 实现细节

TBD

Caffe2

注册OP的代码

Caffe2注册一个OP时候的代码如下:

namespace caffe2 {
namespace {
REGISTER_CPU_OPERATOR(Accumulate, AccumulateOp<float, CPUContext>);

OPERATOR_SCHEMA(Accumulate)
  .NumInputs(1)
  .NumOutputs(1)
  .IdenticalTypeAndShape()
  .SetDoc(R"DOC(
Accumulate operator accumulates the input tensor to the output tensor. If the
output tensor already has the right size, we add to it; otherwise, we first
initialize the output tensor to all zeros, and then do accumulation. Any
further calls to the operator, given that no one else fiddles with the output
in the interim, will do simple accumulations.
Accumulation is done using Axpby operation as shown:
  Y = 1*X + gamma*Y
where X is the input tensor, Y is the output tensor and gamma is the multiplier
argument.
)DOC")
  .Arg("gamma", "(float, default 1.0) Accumulation multiplier")
  .Input(0, "input", "The input tensor that has to be accumulated to the "
         "output tensor. If the output size is not the same as input size, the "
         "output tensor is first reshaped and initialized to zero, and only "
         "then, accumulation is done.")
  .Output(0, "output", "Accumulated output tensor");

SHOULD_NOT_DO_GRADIENT(Accumulate);
}  // namespace
}  // namespace caffe2

这里分为三个部分，他们是:

• 注册前馈网络的计算函数Accumulate
• 注册Accumulate的梯度计算函数，SHOULD_NOT_DO_GRADIENT(Accumulate)。这里很特殊，因为该操作不支持梯度。
• 注册Operator的Schema，也就是操作的元信息。下面分为三个子部分介绍各个模块。

注册forward操作

class OperatorBase {
public:
...
virtual bool Run() = 0;
...
private:
  ...
  vector<const Blob*> inputs_;
  vector<Blob*> outputs_;
  ...
};

template <typename Context>  // Context means device context
class Operator: public OperatorBase {
public
   bool Run() final {
      ...
      RunOnDevice();
      ...
   }
   ...
   virtual bool RunOnDevice() = 0;
   ...
};

template <typename T, typename Context>  // T means data type
class XXXOp final : public Operator<Context> {
public:
	...
	bool RunOnDevice() final {
	  ...
	  T* buf = inputs_[0]->getData<T>();
	  ...
	}
	...
};

首先，Caffe2中，将一层的前馈和反馈操作分离成两个Operation，分别注册，并且都会注册到同一个级别中。即，用户也可以直接调用某一层的反馈函数。

Caffe2中的所有计算OP，都继承自Operator， 而Operator也继承自OperatorBase。他们的作用是

• OperationBase定义了一个OP的对外接口，即所有的Operator只有一个操作Run。而OperationBase中保存了这个操作的protobuf对象(Caffe2基于protobuf2)，输入内存块(const Blob*)和输出内存块(Blob*)的引用。
  • OperationBase并不一定保证Run的时候，所有的Input和Output在同一个设备上。
• Operator定义了一个Operator如何异步执行和跨设备执行。
  • Operator本地保存了这个操作执行的设备信息Context
  • Operator类提供了一个接口RunOnDevice。在RunOnDevice中，保证了Input和Output在统一设备上执行。
  • 统一了Run和RunAsync，都是调用RunOnDevice的函数。
• 每一个Operator直接实现RunOnDevice函数即可，在实现时，使用Input(idx)获取输入，Output(idx)获取输出。将计算结果直接保存在输出中即可。

注册backward操作

这里举两个例子表示Caffe2的反馈操作注册。

REGISTER_CPU_OPERATOR(Accumulate, AccumulateOp<float, CPUContext>);
SHOULD_NOT_DO_GRADIENT(Accumulate);

...

REGISTER_CPU_OPERATOR(FC, FullyConnectedOp<CPUContext>);
REGISTER_CPU_OPERATOR(FCGradient, FullyConnectedGradientOp<CPUContext>);

class GetFCGradient : public GradientMakerBase {
  using GradientMakerBase::GradientMakerBase;
  vector<OperatorDef> GetGradientDefs() override {
    CAFFE_ENFORCE_EQ(def_.input_size(), 3);
    return SingleGradientDef(
        "FCGradient", "",
        vector<string>{I(0), I(1), GO(0)},
        vector<string>{GI(1), GI(2), GI(0)});
  }
};
REGISTER_GRADIENT(FC, GetFCGradient);

首先，注册backward操作，是针对某一个Operator而言的。例如对于Accumulate操作，他的backward操作注册是SHOULD_NOT_DO_GRADIENT(Accumulate)。而对于FC操作，他的backward注册是REGISTER_GRADIENT(FC, GetFCGradient)。

另外，对于一个函数的反馈操作，注册的类型是GradientMakerBase类型。需要注册的函数是vector<OperatorDef> GetGradientDefs()。这里表明了几点问题:

1. 一个Operator的backward操作可能是多个操作。因为返回的是一个Vector。
2. 注册的backward对象直接返回的是protobuf，也就是直接将backward需要的操作append到计算图中。即，在Caffe2执行的过程中，是不区分某一个Op究竟是forward还是backward的。全是配置成同样类型的OperatorDef，再交由执行引擎执行。
3. 在GradientMakerBase中，定义了Backward Operator是和Forward Operator共享了输入和输出。只是可以访问Gradient Buffer。即Gxxx。
   • I(x)表示forward中第几个输入
   • GO(x)表示forward中，第几个输出的Gradient
   • GI(x)表示forward中，第几个输入的Gradient
4. Caffe2不支持变长个数的输入。例如fc_layer，理论上输入是无限多个，但是这种注册机制，只能注册定长的输入个数。

注册元信息

在Caffe2中，元信息被称作Schema。其注册元信息的代码如下:

OPERATOR_SCHEMA(Accumulate)
  .NumInputs(1)
  .NumOutputs(1)
  .IdenticalTypeAndShape()
  .SetDoc(R"DOC(
Accumulate operator accumulates the input tensor to the output tensor. If the
output tensor already has the right size, we add to it; otherwise, we first
initialize the output tensor to all zeros, and then do accumulation. Any
further calls to the operator, given that no one else fiddles with the output
in the interim, will do simple accumulations.
Accumulation is done using Axpby operation as shown:
  Y = 1*X + gamma*Y
where X is the input tensor, Y is the output tensor and gamma is the multiplier
argument.
)DOC")
  .Arg("gamma", "(float, default 1.0) Accumulation multiplier")
  .Input(0, "input", "The input tensor that has to be accumulated to the "
         "output tensor. If the output size is not the same as input size, the "
         "output tensor is first reshaped and initialized to zero, and only "
         "then, accumulation is done.")
  .Output(0, "output", "Accumulated output tensor");

Caffe2的元信息存储在OpSchema对象中，该对象将大部分信息以强类型的形式存储到该对象中。该对象的代码为:

class OpSchema {
private:
  string file_;
  string doc_;
  std::vector<std::pair<const char*, const char*>> arg_desc_{};
  std::vector<std::pair<const char*, const char*>> input_desc_{};
  std::vector<std::pair<const char*, const char*>> output_desc_{};
  int line_ = 0;
  int min_input_ = 0;
  int max_input_ = std::numeric_limits<int>::max();
  int min_output_ = 0;
  int max_output_ = std::numeric_limits<int>::max();
  bool private_ = false;
  std::function<bool(int)> num_inputs_allowed_
      = [](int) { return true; };
  std::function<bool(int)> num_outputs_allowed_
      = [](int) { return true; };
  std::function<bool(int, int)> num_inputs_outputs_allowed_
      = [](int, int) { return true; };
  std::function<int(int)> calculate_output_;
  // In default, any in-place operation is neither allowed nor enforced.
  std::function<bool(int, int)> inplace_allowed_
      = [](int, int) { return false; };
  std::function<bool(int, int)> inplace_enforced_
      = [](int, int) { return false; };
  TensorInferenceFunctionType tensor_inference_function_ =
      [](const OperatorDef& def, const vector<TensorShape>&) {
        vector<TensorShape> out;
        for(int i=0; i<def.output_size(); i++) {
          TensorShape ts;
          ts.set_unknown_shape(true);
          out.push_back(ts);
        }
        return out;
      };
  CostInferenceFunctionType cost_inference_function_ =
      [](const OperatorDef& def, const vector<TensorShape>&) {
        CAFFE_THROW("No cost inference function registered.");
        return Cost();
      };
};

这里有几点需要注意的:

1. 大部分Schema都被定义成了std::function。因为其实Schema的主要做用就是Validation用户的配置。设置成std::function具有很强的灵活性。
2. 这里包含了推测Operator输出Shape的能力。
3. 由于所有类型都是强类型的，增删某一个元信息，或者某一个元信息没有被设置很难表示。未来维护成本可能比较高。
4. 有一些潜在可能发生的不一致性。例如在注释中，写了default value是1.0，但是实际这个default value的是在从protobuf中获取的数据的时候得到的。
5. Arg支持多种数据类型，实际使用protobuf的message搞定。并且获得某一个参数值使用线性查找而不是map。

message Argument {
  optional string name = 1;
  optional float f = 2;
  optional int64 i = 3;
  optional bytes s = 4;
  repeated float floats = 5;
  repeated int64 ints = 6;
  repeated bytes strings = 7;
}

PyTorch

PyTorch的实现与其他神经网络框架有着本质的差别。PyTorch并不试图构造一个high level的接口，也并不试图构造一个『计算图』并执行这个计算图。他的做法相当激进:

• PyTorch其实实现了一个GPU上的numpy和一些神经网络的操作。所有的操作都是即时发生的。

class Pow(Function):
    @staticmethod
    def forward(ctx, a, b):
        ctx.b_size = b.size()
        ctx.save_for_backward(a, b)
        return a.pow(b)

    @staticmethod
    def backward(ctx, grad_output):
        a, b = ctx.saved_variables
        grad_a = grad_output.mul(b).mul(a.pow(b - 1))
        grad_b = grad_output.mul(a.pow(b)).mul(a.log())
        return grad_a, maybe_view(grad_b, ctx.b_size)

• PyTorch借鉴了AutoGrad的实现方式，注册了一些函数如何做backward。
  • 它基本上完成了Caffe2的REGISTER_GRADIENT操作，只是混合了计算过程和注册过程。在forward的时候，它将使用的变量注册到ctx中，而backward中，再读ctx中的内容，进而方便backward

而PyTorch中的所有计算Kernel的实现全是纯C，而且直接实现了Python的接口。所有的参数合法性检查，都写在Kernel函数内部。当然，因为是每个操作都是即时操作的，所以并没有必要去推到输出shape的大小(因为都已经算出来输出是什么了，自然有shape)。

这样实现的缺点也很明显。首先，多线程的实现在PyTorch中是使用openmp和MKL之类的库，对矩阵乘发直接并行，这通信成本应该比较高。而对多显卡并行，PyTorch提供了分发参数和聚合梯度的库，用户自己要去管理什么时刻分发参数，什么时刻聚合参数。

另外，GPU的计算速度应该并不块，因为无论何时何种设备，PyTorch都可以在CPU端读取每一层的输出。这如果写代码不小心的话，很容易占用过多的显存-内存痛心的带宽。

PyTorch是没有计算图的概念的。虽然他在Context中记录了之前调用过的函数，也类似于一种计算图表示。但本质上不是向其他神经网络框架具有两个步骤，设计计算图，使用计算引擎计算计算图。这就很难做性能优化了。

Dynet

DyNet注册OP比较简单，只需要继承自Node类型即可。需要实现的接口包括:

  /**
   * \brief Forward computation
   * \details This function contains the logic for the forward pass. Some implementation remarks from nodes.cc:
   * 1. fx can be understood as a pointer to the (preallocated) location for the result of forward to be stored
   * 2. fx is not initialized, so after calling forward fx must point to the correct answer
   * 3. fx can be repointed to an input, if forward(x) evaluates to x (e.g., in reshaping)
   * 4. scalars results of forward are placed in fx.v[0]
   * 5. DYNET manages its own memory, not Eigen, and it is configured with the EIGEN_NO_MALLOC option. If you get an error about Eigen attempting to allocate memory, it is (probably) because of an implicit creation of a temporary variable. To tell Eigen this is not necessary, the noalias() method is available. If you really do need a temporary variable, its capacity must be requested by Node::aux_storage_size
   *
   * Note on debugging problems with differentiable components
   *
   * - fx is uninitialized when forward is called- are you relying on it being 0?
   *
   * \param xs Pointers to the inputs
   * \param fx pointer to the (preallocated) location for the result of forward to be stored
   */
  virtual void forward_impl(const std::vector<const Tensor*>& xs,
                            Tensor& fx) const = 0;
  //
  /**
   * \brief Accumulates the derivative of E with respect to the ith argument to f, that is, xs[i]
   * \details This function contains the logic for the backward pass. Some implementation remarks from nodes.cc:
   * 1. dEdxi MUST **ACCUMULATE** a result since multiple calls to forward may depend on the same x_i. Even, e.g., Identity must be implemented as dEdx1 += dEdf. THIS IS EXTREMELY IMPORTANT
   * 2. scalars results of forward are placed in fx.v[0]
   * 3. DYNET manages its own memory, not Eigen, and it is configured with the EIGEN_NO_MALLOC option. If you get an error about Eigen attempting to allocate memory, it is (probably) because of an implicit creation of a temporary variable. To tell Eigen this is not necessary, the noalias() method is available. If you really do need a temporary variable, its capacity must be requested by Node::aux_storage_size
   *
   * Note on debugging problems with differentiable components
   *
   * - dEdxi must accummulate (see point 4 above!)
   *
   * \param xs Pointers to inputs
   * \param fx Output
   * \param dEdf Gradient of the objective w.r.t the output of the node
   * \param i Index of the input w.r.t which we take the derivative
   * \param dEdxi Gradient of the objective w.r.t the input of the node
   */
  virtual void backward_impl(const std::vector<const Tensor*>& xs,
                             const Tensor& fx,
                             const Tensor& dEdf,
                             unsigned i,
                             Tensor& dEdxi) const = 0;

这里需要注意的有几点问题:

1. 一个计算单元的前馈和反馈是成对出现的。这里可能会有潜在的代码重用问题。他们是:
   1. 因为一个计算的backward函数可能是由N个forward函数组合而成。而在DyNet的实现中，计算图里面没有backward计算节点，而是将所有节点都用Node表示。
   2. 很难实现backward of backward，难道还要写一个backward_backward_impl?这主要原因也是计算图里面没有backward操作的节点导致的。
2. 一个forward函数与其实多个backward函数相对应。具体使用方法类似于

std::vector<const Tensor*> inputs;
Tensor output;
forward(inputs, output);
Tensor outputGrad;
std::vector<Tensor> inputGrads(inputs.size());

for (size_t i=0; i<outputs.size(); ++i) {
   backward(inputs, output, outputGrad, i, inputGrads[i]);
}

这样做的好处是，可以选择性质的backward部分输入，降低计算量。

DyNet的实现中，另一个比较有趣的地方是，他的计算图结构配置在C++里，使用ComputationGraph对象，而不是使用protobuf之类的东西配置。

• Python也是直接操作ComputationGraph对象。
• 因为大家都直接操作计算图表示的对象，导致计算图的修改实现效率非常高。于是在实际使用中，可以运行时动态的修改计算图，进而这个框架的名字叫DyNet。

Tensorflow

TBD

对比

TBD