手把手教你如何自己設計實現一個深度學習框架（附代碼實現）

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

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

目錄

1. 組件抽象
2. 組件實現
3. 整體結構
4. MNIST 例子
5. 總結
6. 附錄
7. 參考

組件抽象

首先考慮神經網絡運算的流程，神經網絡運算主要包含訓練 training 和預測 predict （或 inference） 兩個階段，訓練的基本流程是：輸入數據 -> 網絡層前向傳播 -> 計算損失 -> 網絡層反向傳播梯度 -> 更新參數，預測的基本流程是 輸入數據 -> 網絡層前向傳播 -> 輸出結果。從運算的角度看，主要可以分為三種類型的計算：

1. 數據在網絡層之間的流動：前向傳播和反向傳播可以看做是張量 Tensor（多維數組）在網絡層之間的流動（前向傳播流動的是輸入輸出，反向傳播流動的是梯度），每個網絡層會進行一定的運算，然后將結果輸入給下一層
2. 計算損失：銜接前向和反向傳播的中間過程，定義了模型的輸出與真實值之間的差異，用來后續(xù)提供反向傳播所需的信息
3. 參數更新：使用計算得到的梯度對網絡參數進行更新的一類計算

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

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

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

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

基本的框架圖如下圖

組件實現

按照上面的抽象，我們可以寫出整個流程代碼如下。

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

首先定義 net，net 的輸入是多個網絡層，然后將 net、loss、optimizer 一起傳給 model。model 實現了 forward、backward 和 apply_grad 三個接口分別對應前向傳播、反向傳播和參數更新三個功能。接下來我們看這里邊各個部分分別如何實現。

tensor

tensor 張量是神經網絡中基本的數據單位，我們這里直接使用 numpy.ndarray?類作為 tensor 類的實現

numpy.ndarray?：https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html

layer

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

#?layer.py
class?Layer(object):
????def?__init__(self,?name):
????????self.name?=?name
????????self.params,?self.grads?=?None,?None

????def?forward(self,?inputs):
????????raise?NotImplementedError

????def?backward(self,?grad):
????????raise?NotImplementedError

最基礎的一種網絡層是全連接網絡層，實現如下。forward 方法接收上層的輸入 inputs，實現 的運算；backward 的方法接收來自上層的梯度，計算關于參數 和輸入的梯度，然后返回關于輸入的梯度。這三個梯度的推導可以見附錄，這里直接給出實現。w_init 和 b_init 分別是參數 和 的初始化器，這個我們在另外的一個實現初始化器中文件 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?=?None

????def?forward(self,?inputs):
????????self.inputs?=?inputs
????????return?inputs?@?self.params["w"]?+?self.params["b"]

????def?backward(self,?grad):
????????self.grads["w"]?=?self.inputs.T?@?grad
????????self.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?=?None

????def?forward(self,?inputs):
????????self.inputs?=?inputs
????????return?self.func(inputs)

????def?backward(self,?grad):
????????return?self.derivative_func(self.inputs)?*?grad

????def?func(self,?x):
????????raise?NotImplementedError

????def?derivative_func(self,?x):
????????raise?NotImplementedError


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

#?net.py
class?Net(object):
????def?__init__(self,?layers):
????????self.layers?=?layers

????def?forward(self,?inputs):
????????for?layer?in?self.layers:
????????????inputs?=?layer.forward(inputs)
????????return?inputs

????def?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.grads

????def?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]

losses

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

梯度稍微復雜一點，目標類別和非目標類別的計算公式不同。對于目標類別維度，其梯度為對應維度模型輸出概率減一，對于非目標類別維度，其梯度為對應維度輸出概率本身。

代碼實現如下

#?loss.py
class?BaseLoss(object):
????def?loss(self,?predicted,?actual):
????????raise?NotImplementedError

????def?grad(self,?predicted,?actual):
????????raise?NotImplementedError


class?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)?/?m

????def?grad(self,?predicted,?actual):
????????m?=?predicted.shape[0]
????????grad?=?np.copy(predicted)
????????grad?-=?actual
????????return?grad?/?m

optimizer

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

#?optimizer.py
class?BaseOptimizer(object):
????def?__init__(self,?lr,?weight_decay):
????????self.lr?=?lr
????????self.weight_decay?=?weight_decay

????def?compute_step(self,?grads,?params):
????????step?=?list()
????????#?flatten?all?gradients
????????flatten_grads?=?np.concatenate(
????????????[np.ravel(v)?for?grad?in?grads?for?v?in?grad.values()])
????????#?compute?step
????????flatten_step?=?self._compute_step(flatten_grads)
????????#?reshape?gradients
????????p?=?0
????????for?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?*?v
????????????????layer[k]?=?_step
????????????????p?+=?block
????????????step.append(layer)
????????return?step

????def?_compute_step(self,?grad):
????????raise?NotImplementedError

class?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,?beta2
????????self._eps?=?eps

????????self._t?=?0
????????self._m,?self._v?=?0,?0

????def?_compute_step(self,?grad):
????????self._t?+=?1
????????self._m?=?self._b1?*?self._m?+?(1?-?self._b1)?*?grad
????????self._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?=?net
????????self.loss?=?loss
????????self.optimizer?=?optimizer

????def?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,?step

????def?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
│   ├── initializer.py
│   ├── layer.py
│   ├── loss.py
│   ├── model.py
│   ├── net.py
│   └── optimizer.py

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

MNIST 例子

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

• 數據集：MNIST（http://yann.lecun.com/exdb/mnist/）
• 任務類型：多分類
• 網絡結構：三層全連接 INPUT(784) -> FC(400) -> FC(100) -> OUTPUT(10)，這個網絡接收 的輸入，其中 是每次輸入的樣本數，784 是每張 的圖像展平后的向量，輸出維度為 ，其中 是樣本數，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):
???????#?training
????????pred?=?model.forward(batch.inputs)
????????loss,?grads?=?model.backward(pred,?batch.targets)
????????model.apply_grad(grads)
????#?evaluate?every?epoch
????test_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 實現了相同的網絡結構、采用相同的采數初始化方法、優(yōu)化器配置等等，得到的結果如下

# 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}

可以看到兩者效果上大差不差，測試集準確率都收斂到 0.982 左右，就單次的實驗看比 Tensorflow 稍微好一點點。

總結

tinynn 相關的源代碼在這個 repo（https://github.com/borgwang/tinynn） 里。目前支持：

• layer ：全連接層、2D 卷積層、 2D反卷積層、MaxPooling 層、Dropout 層、BatchNormalization 層、RNN 層以及 ReLU、Sigmoid、Tanh、LeakyReLU、SoftPlus 等激活函數
• loss：SigmoidCrossEntropy、SoftmaxCrossEntroy、MSE、MAE、Huber
• optimizer：RAam、Adam、SGD、RMSProp、Momentum 等優(yōu)化器，并且增加了動態(tài)調節(jié)學習率 LRScheduler
• 實現了 mnist（分類）、nn_paint（回歸）、DQN（強化學習）、AutoEncoder 和 DCGAN （無監(jiān)督）等常見模型。見 tinynn/examples：https://github.com/borgwang/tinynn/tree/master/examples

tinynn 還有很多可以繼續(xù)完善的地方受限于時間還沒有完成，筆者在空閑時間會進行維護和更新。

當然 tinynn 只是一個「玩具」版本的深度學習框架，一個成熟的深度學習框架至少還需要：支持自動求導、高運算效率（靜態(tài)語言加速、支持 GPU 加速）、提供豐富的算法實現、提供易用的接口和詳細的文檔等等。這個小項目的出發(fā)點更多地是學習，在設計和實現 tinynn 的過程中筆者個人學習確實到了很多東西，包括如何抽象、如何設計組件接口、如何更效率的實現、算法的具體細節(jié)等等。對筆者而言寫這個小框架除了了解深度學習框架的設計與實現之外還有一個好處：后續(xù)可以在這個框架上快速地實現一些新的算法，新的參數初始化方法，新的優(yōu)化算法，新的網絡結構設計，都可以快速地在這個小框架上進行實驗。如果你對自己設計實現一個深度學習框架也感興趣，希望看完這篇文章會對你有所幫助，也歡迎大家提 PR 一起貢獻代碼~

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

多分類下交叉熵損失如下式：

其中 分別是真實值和模型預測值， 是樣本數， 是類別個數。由于真實值一般為一個 one-hot 向量（除了真實類別維度為 1 其他均為 0），因此上式可以化簡為

其中 是代表真實類別， 代表第 個樣本 類的預測概率。即我們需要計算的是每個樣本在真實類別上的預測概率的對數的和，然后再取負就是交叉熵損失。接下來推導如何求解該損失關于模型輸出的梯度，用 表示模型輸出，在多分類中通常最后會使用 Softmax 將網絡的輸出歸一化為一個概率分布，則 Softmax 后的輸出為

代入上面的損失函數

求解 關于輸出向量 的梯度，可以將 分為目標類別所在維度 和非目標類別維度 。首先看目標類別所在維度
再看非目標類別所在維度

可以看到對于目標類別維度，其梯度為對應維度模型輸出概率減一，對于非目標類別維度，其梯度為對應維度輸出概率真身。

參考

• Deep Learning, Goodfellow, et al. (2016)
• Joel Grus - Livecoding Madness - Let's Build a Deep Learning Library
• TensorFlow Documentation
• PyTorch Documentation

如果覺得有用，就請分享到朋友圈吧！