MXNet 源码分析-基本流程

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

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
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 输入的方法.