從頭搭建一個深度學習框架

轉自：Build a Deep Learning Framework From Scratch

代碼：https://github.com/borgwang/tinynn

在這里插入圖片描述

當前深度學習框架越來越成熟，對于使用者而言封裝程度越來越高，好處就是現在可以非常快速地將這些框架作為工具使用，用非常少的代碼就可以進行實驗，壞處就是可能背后地實現都被隱藏起來了。在這篇文章里筆者將帶大家一起從頭設計和實現一個輕量級的（大約 200 行）、易于擴展的深度學習框架 tinynn，希望對大家了解深度學習框架的基本設計和實現有一定的幫助。

本文首先會從深度學習的流程開始分析，對神經網絡中的關鍵組件抽象，確定基本框架；然后再對框架里各個組件進行代碼實現；最后基于這個框架實現了一個 MNIST 分類的示例。

組件抽象

首先考慮神經網絡運算的流程，神經網絡運算主要包含訓練 training 和預測 predict （或 inference）兩個階段，

訓練的基本流程是：輸入數據 -> 網絡層前向傳播 -> 計算損失 -> 網絡層反向傳播梯度 -> 更新參數；
預測的基本流程是輸入數據 -> 網絡層前向傳播 -> 輸出結果。

從運算的角度看，主要可以分為三種類型的計算：

數據在網絡層直接的流動

前向傳播和反向傳播可以看做是張量 Tensor（多維數組）在網絡層之間的流動（前向傳播流動的是輸入輸出，反向傳播流動的是梯度），每個網絡層會進行一定的運算，然后將結果輸入給下一層。
計算損失

銜接前向和反向傳播的中間過程，定義了模型的輸出與真實值之間的差異，用來后續提供反向傳播所需的信息
參數更新

使用計算得到的梯度對網絡參數進行更新的一類計算

基于這個三種類型，我們可以對網絡的基本組件做一個抽象

tensor 張量，這個是神經網絡中數據的基本單位
layer 網絡層，負責接收上一層的輸入，進行該層的運算，將結果輸出給下一層，由于 tensor 的流動有前向和反向兩個方向，因此對于每種類型網絡層我們都需要同時實現 forward 和 backward 兩種運算
loss 損失，在給定模型預測值與真實值之后，該組件輸出損失值以及關于最后一層的梯度（用于梯度回傳）
optimizer 優化器，負責使用梯度更新模型的參數

然后我們還需要一些組件把上面這個 4 種基本組件整合到一起，形成一個 pipeline

net 組件負責管理 tensor 在 layer 之間的前向和反向傳播，同時能提供獲取參數、設置參數、獲取梯度的接口
model 組件負責整合所有組件，形成整個 pipeline。即 net 組件進行前向傳播 -> loss 組件計算損失和梯度 -> net 組件將梯度反向傳播 -> optimizer 組件將梯度更新到參數。

基本框架如下圖

在這里插入圖片描述

組件實現

按照上面的抽象，我們可以寫出整個流程代碼如下。首先定義 net，net 的輸入是多個網絡層，然后將 net、loss、optimizer 一起傳給 model。model 實現了 forward、backward 和 apply_grad 三個接口分別對應前向傳播、反向傳播和參數更新三個功能。

# define model
net = Net([layer1, layer2, ...])
model = Model(net, loss_fn, optimizer)# training
pred = model.forward(train_X)
loss, grads = model.backward(pred, train_Y)
model.apply_grad(grads)# inference
test_pred = model.forward(test_X)

接下來我們看這里邊各個部分分別如何實現。

tensor

tensor 張量是神經網絡中基本的數據單位，我們這里直接使用 numpy.ndarray 類作為 tensor 類的實現（numpy 底層使用了 C 和 Fortran，并且在算法層面進行了大量的優化，運算速度也不算慢）

layer

上面流程代碼中 model 進行 forward 和 backward，其實底層都是網絡層在進行實際運算，因此網絡層需要有提供 forward 和 backward 接口進行對應的運算。同時還應該將該層的參數和梯度記錄下來。先實現一個基類如下

# layer.py
class Layer(object):def __init__(self, name):self.name = nameself.params, self.grads = None, Nonedef forward(self, inputs):raise NotImplementedErrordef backward(self, grad):raise NotImplementedError

最基礎的一種網絡層是全連接網絡層，實現如下。forward 方法接收上層的輸入 inputs，實現 $w x + b$ 的運算；backward 的方法接收來自上層的梯度，計算關于參數 $w, b$ 和輸入的梯度，然后返回關于輸入的梯度。這三個梯度的推導可以見附錄，這里直接給出實現。w_init 和 b_init 分別是參數 weight 和 bias 的初始化器，這個我們在另外的一個實現初始化器中文件 initializer.py 去實現，這部分不是核心部件，所以在這里不展開介紹。

# layer.py
class Dense(Layer):def __init__(self, num_in, num_out,w_init=XavierUniformInit(),b_init=ZerosInit()):super().__init__("Linear")self.params = {"w": w_init([num_in, num_out]),"b": b_init([1, num_out])}self.inputs = Nonedef forward(self, inputs):self.inputs = inputsreturn inputs @ self.params["w"] + self.params["b"]def backward(self, grad):self.grads["w"] = self.inputs.T @ gradself.grads["b"] = np.sum(grad, axis=0)return grad @ self.params["w"].T

同時神經網絡中的另一個重要的部分是激活函數。激活函數可以看做是一種網絡層，同樣需要實現 forward 和 backward 方法。我們通過繼承 Layer 類實現激活函數類，這里實現了最常用的 ReLU 激活函數。func 和 derivation_func 方法分別實現對應激活函數的正向計算和梯度計算。

# layer.py
class Activation(Layer):"""Base activation layer"""def __init__(self, name):super().__init__(name)self.inputs = Nonedef forward(self, inputs):self.inputs = inputsreturn self.func(inputs)def backward(self, grad):return self.derivative_func(self.inputs) * graddef func(self, x):raise NotImplementedErrordef derivative_func(self, x):raise NotImplementedErrorclass ReLU(Activation):"""ReLU activation function"""def __init__(self):super().__init__("ReLU")def func(self, x):return np.maximum(x, 0.0)def derivative_func(self, x):return x > 0.0

net

上文提到 net 類負責管理 tensor 在 layer 之間的前向和反向傳播。forward 方法很簡單，按順序遍歷所有層，每層計算的輸出作為下一層的輸入；backward 則逆序遍歷所有層，將每層的梯度作為下一層的輸入。這里我們還將每個網絡層參數的梯度保存下來返回，后面參數更新需要用到。另外 net 類還實現了獲取參數、設置參數、獲取梯度的接口，也是后面參數更新時需要用到

# net.py
class Net(object):def __init__(self, layers):self.layers = layersdef forward(self, inputs):for layer in self.layers:inputs = layer.forward(inputs)return inputsdef backward(self, grad):all_grads = []for layer in reversed(self.layers):grad = layer.backward(grad)all_grads.append(layer.grads)return all_grads[::-1]def get_params_and_grads(self):for layer in self.layers:yield layer.params, layer.gradsdef get_parameters(self):return [layer.params for layer in self.layers]def set_parameters(self, params):for i, layer in enumerate(self.layers):for key in layer.params.keys():layer.params[key] = params[i][key]

loss

上文我們提到 loss 組件需要做兩件事情，給定了預測值和真實值，需要計算損失值和關于預測值的梯度。我們分別實現為 loss 和 grad 兩個方法，這里我們實現多分類回歸常用的 SoftmaxCrossEntropyLoss 損失。這個的損失 loss 和梯度 grad 的計算公式推導進文末附錄，這里直接給出結果：

多分類 softmax 交叉熵的損失為
$JCE(y,y^)=?∑i=1Nlogy^icJ_{CE}(y,\hat{y})=-\sum_{i=1}^Nlog\ \hat{y}_i^c$
梯度稍微復雜一點，目標類別和非目標類別的計算公式不同。對于目標類別維度，其梯度為對應維度模型輸出概率減一，對于非目標類別維度，其梯度為對應維度輸出概率本身。
$?JCE?oc={(y^c?1)/N目標類別cyc~/N非目標類別c~\frac{\partial{J_{CE}}}{\partial{o^c}}=\begin{cases} (\hat{y}^c-1)/N\ \ \ \ \ 目標類別c\\ y^{\tilde{c}}/N\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ 非目標類別\tilde{c}\end{cases}$
代碼實現如下
```
# loss.py
class BaseLoss(object):def loss(self, predicted, actual):raise NotImplementedErrordef grad(self, predicted, actual):raise NotImplementedErrorclass CrossEntropyLoss(BaseLoss):def loss(self, predicted, actual):m = predicted.shape[0]exps = np.exp(predicted - np.max(predicted, axis=1, keepdims=True))p = exps / np.sum(exps, axis=1, keepdims=True)nll = -np.log(np.sum(p * actual, axis=1))return np.sum(nll) / mdef grad(self, predicted, actual):m = predicted.shape[0]grad = np.copy(predicted)grad -= actualreturn grad / m
```

optimizer

optimizer 主要實現一個接口 compute_step，這個方法根據當前的梯度，計算返回實際優化時每個參數改變的步長。我們在這里實現常用的 Adam 優化器。

# optimizer.py
class BaseOptimizer(object):def __init__(self, lr, weight_decay):self.lr = lrself.weight_decay = weight_decaydef compute_step(self, grads, params):step = list()# flatten all gradientsflatten_grads = np.concatenate([np.ravel(v) for grad in grads for v in grad.values()])# compute stepflatten_step = self._compute_step(flatten_grads)# reshape gradientsp = 0for param in params:layer = dict()for k, v in param.items():block = np.prod(v.shape)_step = flatten_step[p:p+block].reshape(v.shape)_step -= self.weight_decay * vlayer[k] = _stepp += blockstep.append(layer)return stepdef _compute_step(self, grad):raise NotImplementedErrorclass Adam(BaseOptimizer):def __init__(self, lr=0.001, beta1=0.9, beta2=0.999,eps=1e-8, weight_decay=0.0):super().__init__(lr, weight_decay)self._b1, self._b2 = beta1, beta2self._eps = epsself._t = 0self._m, self._v = 0, 0def _compute_step(self, grad):self._t += 1self._m = self._b1 * self._m + (1 - self._b1) * gradself._v = self._b2 * self._v + (1 - self._b2) * (grad ** 2)# bias correction_m = self._m / (1 - self._b1 ** self._t)_v = self._v / (1 - self._b2 ** self._t)return -self.lr * _m / (_v ** 0.5 + self._eps)

model

最后 model 類實現了我們一開始設計的三個接口 forward、backward 和 apply_grad ，forward 直接調用 net 的 forward ，backward 中把 net 、loss、optimizer 串起來，先計算損失 loss，然后反向傳播得到梯度，然后 optimizer 計算步長，最后由 apply_grad 對參數進行更新

# model.py
class Model(object):def __init__(self, net, loss, optimizer):self.net = netself.loss = lossself.optimizer = optimizerdef forward(self, inputs):return self.net.forward(inputs)def backward(self, preds, targets):loss = self.loss.loss(preds, targets)grad = self.loss.grad(preds, targets)grads = self.net.backward(grad)params = self.net.get_parameters()step = self.optimizer.compute_step(grads, params)return loss, stepdef apply_grad(self, grads):for grad, (param, _) in zip(grads, self.net.get_params_and_grads()):for k, v in param.items():param[k] += grad[k]

整體結構

最后我們實現出來核心代碼部分文件結構如下：

tinynn
├── core
│   ├── __init__.py
│   ├── initializer.py
│   ├── layer.py
│   ├── loss.py
│   ├── model.py
│   ├── net.py
│   └── optimizer.py

其中 initializer.py 這個模塊上面沒有展開講，主要實現了常見的參數初始化方法，用于給網絡層初始化參數。

MNIST例子

框架基本搭起來后，我們找一個例子來用 tinynn 這個框架 run 起來。這個例子的基本一些配置如下

數據集：MNIST
任務類型：多分類
網絡結構：三層全連接 INPUT(784) -> FC(400) -> FC(100) -> OUTPUT(10)，這個網絡接收 $(N, 784)$ 的輸入，其中 $N$ 是每次輸入的樣本數， $784$ 是每張 $(28, 28)$ 的圖像展平后的向量，輸出維度為 $(N, 10)$ ，其中 $N$ 是樣本數， $10$ 是對應圖片在 $10$ 個類別上的概率
激活函數：ReLU
損失函數：SoftmaxCrossEntropy
optimizer：Adam(lr=1e-3)
batch_size：128
Num_epochs：20

這里我們忽略數據載入、預處理等一些準備代碼，只把核心的網絡結構定義和訓練代碼貼出來如下

# example/mnist/run.py
net = Net([Dense(784, 400),ReLU(),Dense(400, 100),ReLU(),Dense(100, 10)
])
model = Model(net=net, loss=SoftmaxCrossEntropyLoss(), optimizer=Adam(lr=args.lr))iterator = BatchIterator(batch_size=args.batch_size)
evaluator = AccEvaluator()
for epoch in range(num_ep):for batch in iterator(train_x, train_y):# trainingpred = model.forward(batch.inputs)loss, grads = model.backward(pred, batch.targets)model.apply_grad(grads)# evaluate every epochtest_pred = model.forward(test_x)test_pred_idx = np.argmax(test_pred, axis=1)test_y_idx = np.asarray(test_y)res = evaluator.evaluate(test_pred_idx, test_y_idx)print(res)

運行結果如下

# tinynn
Epoch 0 	 {'total_num': 10000, 'hit_num': 9658, 'accuracy': 0.9658}
Epoch 1 	 {'total_num': 10000, 'hit_num': 9740, 'accuracy': 0.974}
Epoch 2 	 {'total_num': 10000, 'hit_num': 9783, 'accuracy': 0.9783}
Epoch 3 	 {'total_num': 10000, 'hit_num': 9799, 'accuracy': 0.9799}
Epoch 4 	 {'total_num': 10000, 'hit_num': 9805, 'accuracy': 0.9805}
Epoch 5 	 {'total_num': 10000, 'hit_num': 9826, 'accuracy': 0.9826}
Epoch 6 	 {'total_num': 10000, 'hit_num': 9823, 'accuracy': 0.9823}
Epoch 7 	 {'total_num': 10000, 'hit_num': 9819, 'accuracy': 0.9819}
Epoch 8 	 {'total_num': 10000, 'hit_num': 9820, 'accuracy': 0.982}
Epoch 9 	 {'total_num': 10000, 'hit_num': 9838, 'accuracy': 0.9838}
Epoch 10 	 {'total_num': 10000, 'hit_num': 9825, 'accuracy': 0.9825}
Epoch 11 	 {'total_num': 10000, 'hit_num': 9810, 'accuracy': 0.981}
Epoch 12 	 {'total_num': 10000, 'hit_num': 9845, 'accuracy': 0.9845}
Epoch 13 	 {'total_num': 10000, 'hit_num': 9845, 'accuracy': 0.9845}
Epoch 14 	 {'total_num': 10000, 'hit_num': 9835, 'accuracy': 0.9835}
Epoch 15 	 {'total_num': 10000, 'hit_num': 9817, 'accuracy': 0.9817}
Epoch 16 	 {'total_num': 10000, 'hit_num': 9815, 'accuracy': 0.9815}
Epoch 17 	 {'total_num': 10000, 'hit_num': 9835, 'accuracy': 0.9835}
Epoch 18 	 {'total_num': 10000, 'hit_num': 9826, 'accuracy': 0.9826}
Epoch 19 	 {'total_num': 10000, 'hit_num': 9819, 'accuracy': 0.9819}

可以看到測試集 accuracy 隨著訓練進行在慢慢提升，這說明數據在框架中確實按照正確的方式進行流動和計算。為了對比下效果，我用 Tensorflow (1.13.1) 實現了相同的網絡結構、采用相同的采數初始化方法、優化器配置等等，得到的結果如下

# Tensorflow 1.13.1
Epoch 0 	 {'total_num': 10000, 'hit_num': 9591, 'accuracy': 0.9591}
Epoch 1 	 {'total_num': 10000, 'hit_num': 9734, 'accuracy': 0.9734}
Epoch 2 	 {'total_num': 10000, 'hit_num': 9706, 'accuracy': 0.9706}
Epoch 3 	 {'total_num': 10000, 'hit_num': 9756, 'accuracy': 0.9756}
Epoch 4 	 {'total_num': 10000, 'hit_num': 9722, 'accuracy': 0.9722}
Epoch 5 	 {'total_num': 10000, 'hit_num': 9772, 'accuracy': 0.9772}
Epoch 6 	 {'total_num': 10000, 'hit_num': 9774, 'accuracy': 0.9774}
Epoch 7 	 {'total_num': 10000, 'hit_num': 9789, 'accuracy': 0.9789}
Epoch 8 	 {'total_num': 10000, 'hit_num': 9766, 'accuracy': 0.9766}
Epoch 9 	 {'total_num': 10000, 'hit_num': 9763, 'accuracy': 0.9763}
Epoch 10 	 {'total_num': 10000, 'hit_num': 9791, 'accuracy': 0.9791}
Epoch 11 	 {'total_num': 10000, 'hit_num': 9773, 'accuracy': 0.9773}
Epoch 12 	 {'total_num': 10000, 'hit_num': 9804, 'accuracy': 0.9804}
Epoch 13 	 {'total_num': 10000, 'hit_num': 9782, 'accuracy': 0.9782}
Epoch 14 	 {'total_num': 10000, 'hit_num': 9800, 'accuracy': 0.98}
Epoch 15 	 {'total_num': 10000, 'hit_num': 9837, 'accuracy': 0.9837}
Epoch 16 	 {'total_num': 10000, 'hit_num': 9811, 'accuracy': 0.9811}
Epoch 17 	 {'total_num': 10000, 'hit_num': 9793, 'accuracy': 0.9793}
Epoch 18 	 {'total_num': 10000, 'hit_num': 9818, 'accuracy': 0.9818}
Epoch 19 	 {'total_num': 10000, 'hit_num': 9811, 'accuracy': 0.9811}

在這里插入圖片描述

可以看到兩者效果上大差不差，就單次的實驗看比 Tensorflow 稍微好一點點。

總結

tinynn 相關的源代碼在這個 repo 里。目前支持：

layer：全連接層、2D 卷積層、 2D 反卷積層、MaxPooling 層、Dropout 層、BatchNormalization 層、RNN 層以及 ReLU、Sigmoid、Tanh、LeakyReLU、SoftPlus 等激活函數
loss：SigmoidCrossEntropy、SoftmaxCrossEntroy、MSE、MAE、Huber
optimizer：RAam、Adam、SGD、RMSProp、Momentum 等優化器，并且增加了動態調節學習率 LRScheduler
實現了 mnist（分類）、nn_paint（回歸）、DQN（強化學習）、AutoEncoder 和 DCGAN （無監督）等常見模型。見 tinynn/examples

tinynn 還有很多可以繼續完善的地方受限于時間還沒有完成（實現循環神經網絡層、BatchNorm 層、對運算效率進行優化等等），筆者在空閑時間會進行維護和更新。

當然從生產應用的角度 tinynn 可能是一個糟糕的選擇，理由用 python 實現在這種計算密集型的場景中會不可避免地出現性能問題、沒有 GPU 支持、沒有分布式支持、很多算法還沒實現等等等，這個小項目的出發點更多地是學習，在設計和實現 tinynn 的過程中筆者個人學習確實到了很多東西，包括如何抽象、如何設計組件接口、如何更效率的實現、算法的具體細節等等。對筆者而言這個寫這個小框架除了了解深度學習框架的設計與實現之外還有一個好處：后續可以在這個框架上快速地實現一些新的算法，比如新出了某篇 paper 提出來新的參數初始化方法，新的優化算法，新的網絡結構設計，都可以快速地在這個小框架上實驗。

如果你對自己設計實現一個深度學習框架也感興趣，希望看完這篇文章會對你有所幫助，也歡迎大家提 PR 一起貢獻代碼~

附錄 Softmax 交叉熵損失和梯度推導

多分類下交叉熵損失如下式：
$JCE(y,y^)=?∑i=1N∑k=1Kyiklogyik^J_{CE}(y,\hat{y})=-\sum_{i=1}^N\sum_{k=1}^Ky_i^klog\ \hat{y_i^k}$
其中 $y,y^y,\hat{y}$ 分別是真實值和模型預測值， $N$ 是樣本數， $K$ 是類別個數。由于真實值一般為一個 one-hot 向量（除了真實類別維度為 1 其他均為 0），因此上式可以化簡為
$JCE(y,y^)=?∑i=1Nlogyic^J_{CE}(y,\hat{y})=-\sum_{i=1}^Nlog\ \hat{y_i^c}$
其中 $c$ 是代表真實類別， $yic^\hat{y_i^c}$ 代表第 $i$ 個樣本 $c$ 類的預測概率。即我們需要計算的是每個樣本在真實類別上的預測概率的對數的和，然后再取負就是交叉熵損失。

接下來推導如何求解該損失關于模型輸出的梯度，用 $o$ 表示模型輸出，在多分類中通常最后會使用 Softmax 將網絡的輸出歸一化為一個概率分布，則 Softmax 后的輸出為
$yc^=exp(oc)∑k=1Kexp(ok)\hat{y_c}=\frac{exp(o^c)}{\sum_{k=1}^Kexp(o^k)}$
代入上面的損失函數
$JCE=?∑i=1N(oic?log∑k=1Kexp(oik))J_{CE}=-\sum_{i=1}^N(o_i^c-log\sum_{k=1}^Kexp(o_i^k))$
求解 $J_{CE}$ 關于輸出向量 $o$ 的梯度，可以將 $o$ 分為目標類別所在維度 $o^c$ 和非目標類別維度 $KaTeX parse error: Got function '\tilde' with no arguments as superscript at position 3: o^\?t?i?l?d?e?{c}$ 。首先看目標類別所在維度 $o^c$
$?JCE?oc=?∑i=1N(1?exp(oc)∑k=1Kexp(ok))=∑i=1N(y^?1)\frac{\partial{J_{CE}}}{\partial{o^c}}=-\sum_{i=1}^N(1-\frac{exp(o^c)}{\sum_{k=1}^Kexp(o^k)})=\sum_{i=1}^N(\hat{y}-1)$
再看非目標類別所在維度 $oc~o^{\tilde{c}}$
$?JCE?oc~=?∑i=1N(?exp(oc)∑k=1Kexp(ok))=∑i=1N(y^)\frac{\partial{J_{CE}}}{\partial{o^{\tilde{c}}}}=-\sum_{i=1}^N(-\frac{exp(o^c)}{\sum_{k=1}^Kexp(o^k)})=\sum_{i=1}^N(\hat{y})$
可以看到對于目標類別維度，其梯度為對應維度模型輸出概率減一，對于非目標類別維度，其梯度為對應維度輸出概率本身。