MXNet 源码分析-基本流程

对于用户来说，MXNet 工作的基本流程是，首先通过 symbol 定义一个 computation graph, 然后，把数据绑定到 symbol 上，最后，执行 forward 和 backward 完成模型的优化。下面通过一个最简单的例子来说明这个过程。

C++ API 提供了一个简单的例子：

void MLP() {
  auto sym_x = Symbol::Variable("X");  //定义一个 Symbol, 网络的出入节点，后面会把网络的输入数据绑定到该 Symbol(内部是绑定的 computation graph 的 node 上去)
  auto sym_label = Symbol::Variable("label");  // 同上，输入数据的标签信息

  const int nLayers = 2;
  vector<int> layerSizes({512, 10});
  vector<Symbol> weights(nLayers);  // 该例子中的 Symbol 是 cpp 接口中定义的 Symbol, 封装了 MXNet 的 Symbol
  vector<Symbol> biases(nLayers);
  vector<Symbol> outputs(nLayers);

  // 定义网络结构，各个具体的 Symbol 的定义可参考 cpp 接口。
  for (int i = 0; i < nLayers; i++) {
    string istr = to_string(i);
    weights[i] = Symbol::Variable(string("w") + istr);
    biases[i] = Symbol::Variable(string("b") + istr);
    Symbol fc = FullyConnected(string("fc") + istr,
      i == 0? sym_x : outputs[i-1],
      weights[i], biases[i], layerSizes[i]);
    outputs[i] = LeakyReLU(string("act") + istr, fc, LeakyReLUActType::kLeaky);
  }
  auto sym_out = SoftmaxOutput("softmax", outputs[nLayers - 1], sym_label);

  // 定义该模型要在哪个设备上计算，CPU, 哪块 GPU. CPU 默认是使用所有的 GPU 的所有核心，
  // GPU 默认在第 0 块显卡计算，如果选择其它显卡需要具体指定.(并不会默认使用所有的显卡)
  Context ctx_dev(DeviceType::kCPU, 0);

  //定义输入数据，并且分配相应的 memory, 如果第三个参数是 true 的话，执行 delay allocation, 不会立刻分配 memory
  NDArray array_x(Shape(128, 28), ctx_dev, false);
  NDArray array_y(Shape(128), ctx_dev, false);

  mx_float* aptr_x = new mx_float[128 * 28];
  mx_float* aptr_y = new mx_float[128];

  // we make the data by hand, in 10 classes, with some pattern
  for (int i = 0; i < 128; i++) {
    for (int j = 0; j < 28; j++) {
      aptr_x[i * 28 + j] = i % 10 * 1.0f;
    }
    aptr_y[i] = i % 10;
  }

  // 把具体的数据同步到 NDArray 中，之所以有这个显示的方法，主要是考虑到使用 GPU 计算的使用，需要首先把数据同步到显存上去
  array_x.SyncCopyFromCPU(aptr_x, 128 * 28);
  // 等待数据同步完成
  array_x.WaitToRead();
  array_y.SyncCopyFromCPU(aptr_y, 128);
  array_y.WaitToRead();

  // 定义网络的参数  NDArray
  NDArray array_w_1(Shape(512, 28), ctx_dev, false);
  NDArray array_b_1(Shape(512), ctx_dev, false);
  NDArray array_w_2(Shape(10, 512), ctx_dev, false);
  NDArray array_b_2(Shape(10), ctx_dev, false);

  // 初始化参数
  array_w_1 = 0.5f;
  array_b_1 = 0.0f;
  array_w_2 = 0.5f;
  array_b_2 = 0.0f;

  // 定义存放梯度的 NDArray
  NDArray array_w_1_g(Shape(512, 28), ctx_dev, false);
  NDArray array_b_1_g(Shape(512), ctx_dev, false);
  NDArray array_w_2_g(Shape(10, 512), ctx_dev, false);
  NDArray array_b_2_g(Shape(10), ctx_dev, false);

  // 把网络输入和参数 (可以统一看做该 computation graph 计算过程中需要的数据) 放到一个 std::vector 中 (为了后面的构建 executor)
  std::vector<NDArray> in_args;
  in_args.push_back(array_x);
  in_args.push_back(array_w_1);
  in_args.push_back(array_b_1);
  in_args.push_back(array_w_2);
  in_args.push_back(array_b_2);
  in_args.push_back(array_y);
  //  与上面基本相同的意思，注意，顺序要和上面的 in_args 相同，即使有些参数不需要计算梯度，也要 push_back 一个空的 NDArray
  std::vector<NDArray> arg_grad_store;
  arg_grad_store.push_back(NDArray());
  arg_grad_store.push_back(array_w_1_g);
  arg_grad_store.push_back(array_b_1_g);
  arg_grad_store.push_back(array_w_2_g);
  arg_grad_store.push_back(array_b_2_g);
  arg_grad_store.push_back(NDArray());
  // 指定梯度 memory 的更新方法，顺序同样要和上面的严格一致。kNullOp 表示没有任何操作，kWriteTo 表示用新计算出来的值覆盖原来的值
  std::vector<OpReqType> grad_req_type;
  grad_req_type.push_back(kNullOp);
  grad_req_type.push_back(kWriteTo);
  grad_req_type.push_back(kWriteTo);
  grad_req_type.push_back(kWriteTo);
  grad_req_type.push_back(kWriteTo);
  grad_req_type.push_back(kNullOp);

  // 辅助状态，这里没有用到，如果是有 batch normalization 的话，该参数就需要像上面一样具体定义了
  std::vector<NDArray> aux_states;

  cout << "make the Executor" << endl;

  // 构建 executor, 该步骤执行了非常多的操作，具体讲至少有以下几个重要操作：
  // 1. 检查参数的上面的参数 NDArray 的 shape
  // 2. 构建 computation graph
  // 3. 构建具体的 Engine(同步还是异步，是否使用线程池等)
  // 4. 推导 graph 的 topo 序
  // 5. 把 NDArray 参数按照图谱序绑定到具体的 node 上
  Executor* exe = new Executor(sym_out, ctx_dev, in_args, arg_grad_store,
                               grad_req_type, aux_states);

  cout << "Training" << endl;
  int max_iters = 20000;
  mx_float learning_rate = 0.0001;
  for (int iter = 0; iter < max_iters; ++iter) {
    exe->Forward(true);  //前向计算

    if (iter % 100 == 0) {
      cout << "epoch " << iter << endl;
      std::vector<NDArray>& out = exe->outputs;
      float* cptr = new float[128 * 10];
      out[0].SyncCopyToCPU(cptr, 128 * 10);
      NDArray::WaitAll();
      OutputAccuracy(cptr, aptr_y);
      delete[] cptr;
    }

    // update the parameters
    exe->Backward(); //反向计算梯度
    for (int i = 1; i < 5; ++i) {
      in_args[i] -= arg_grad_store[i] * learning_rate; //更新参数
    }
    NDArray::WaitAll();
  }

  delete exe;
  delete[] aptr_x;
  delete[] aptr_y;
}

上面代码需要重点注意的点：

1. 拓扑序非常重要，拓扑序中，对于某个节点 N, 该节点要么没有其它依赖，要么仅仅依赖于该节点前面的节点，因此，computation graph 进行拓扑排序之后，直接按照 topo 序进行计算就可以了.(异步模式另行讨论)
2. 上面的代码中各个输入输入在 push_back 到 std::vector 中的时候其实是作者已经知道了该网络在转换成 computation graph 之后的拓扑序，然后按照该拓扑序进行 push_back 的
3. MXNet 的 cpp 借口和 Python 借口都提供了 simple_bind 方法，该方法封装了上面 NDArray 参数 push_back 到 std::vector 的过程，其大概过程同样是先对 computation graph 进行拓扑排序，然后，获取排序后的各个节点需要的参数的顺序。因为在每个 Symbol 中都有 name 的，因此，simple_bind 要求输入的 NDArray 也要有 name, 而且，NDArray 中的 name 要和 Symbol 中的 name 对应起来。这样，用户只需要保证 NDArray 的 name 能和 Symbol 中的 name 对应就可以了，而不用保证其 topo 序是 match 的。
4. 上述代码每次都是使用相同的数据 forward, 也就是相当于 Gradient Descend 方法，mini batch 的 SGD 要求每次输入一个 batch 的数据。那么，MXNet 是怎么每次输入一个 batch 的数据的呢？每次都重新 new 一个 Executor 吗？毕竟所有的数据都是在 new Executor 的时候传入的。显然这样的方法会比较蠢，MXNet 当然提供了 feed 输入的方法。