手把手教你如何自己設計實現一個深度學習框架(附程式碼實現)


MLNLP

社群是國內外知名的機器學習與自然語言處理社群,受眾覆蓋國內外NLP碩博生、高校老師以及企業研究人員。


社群的願景是促進國內外自然語言處理,機器學習學術界、產業界和廣大愛好者之間的交流和進步,特別是初學者同學們的進步。
轉載自 | 極市平臺
作者丨王桂波@知乎
來源丨https://zhuanlan.zhihu.com/p/78713744
當前深度學習框架越來越成熟,對於使用者而言封裝程度越來越高,好處就是現在可以非常快速地將這些框架作為工具使用,用非常少的程式碼就可以構建模型進行實驗,壞處就是可能背後地實現都被隱藏起來了。在這篇文章裡筆者將設計和實現一個、輕量級的(約 200 行)、易於擴充套件的深度學習框架 tinynn(基於 Python 和 Numpy 實現),希望對大家瞭解深度學習的基本元件、框架的設計和實現有一定的幫助。
本文首先會從深度學習的流程開始分析,對神經網路中的關鍵元件抽象,確定基本框架;然後再對框架裡各個元件進行程式碼實現;最後基於這個框架實現了一個 MNIST 分類的示例,並與 Tensorflow 做了簡單的對比驗證。
1
『元件抽象』
首先考慮神經網路運算的流程,神經網路運算主要包含訓練 training 和預測 predict (或 inference) 兩個階段,訓練的基本流程是:輸入資料 -> 網路層前向傳播 -> 計算損失 -> 網路層反向傳播梯度 -> 更新引數,預測的基本流程是 輸入資料 -> 網路層前向傳播 -> 輸出結果。從運算的角度看,主要可以分為三種類型的計算:
  1. 資料在網路層之間的流動:前向傳播和反向傳播可以看做是張量 Tensor(多維陣列)在網路層之間的流動(前向傳播流動的是輸入輸出,反向傳播流動的是梯度),每個網路層會進行一定的運算,然後將結果輸入給下一層
  2. 計算損失:銜接前向和反向傳播的中間過程,定義了模型的輸出與真實值之間的差異,用來後續提供反向傳播所需的資訊
  3. 引數更新:使用計算得到的梯度對網路引數進行更新的一類計算
基於這個三種類型,我們可以對網路的基本元件做一個抽象
  • tensor 張量,這個是神經網路中資料的基本單位
  • layer 網路層,負責接收上一層的輸入,進行該層的運算,將結果輸出給下一層,由於 tensor 的流動有前向和反向兩個方向,因此對於每種型別網路層我們都需要同時實現 forward 和 backward 兩種運算
  • loss 損失,在給定模型預測值與真實值之後,該元件輸出損失值以及關於最後一層的梯度(用於梯度回傳)
  • optimizer 最佳化器,負責使用梯度更新模型的引數
然後我們還需要一些元件把上面這個 4 種基本元件整合到一起,形成一個 pipeline
  • net 元件負責管理 tensor 在 layers 之間的前向和反向傳播,同時能提供獲取引數、設定引數、獲取梯度的介面
  • model 元件負責整合所有元件,形成整個 pipeline。即 net 元件進行前向傳播 -> losses 元件計算損失和梯度 -> net 元件將梯度反向傳播 -> optimizer 元件將梯度更新到引數。
基本的框架圖如下圖
2
『元件實現』
按照上面的抽象,我們可以寫出整個流程程式碼如下。

# 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
classLayer(object):
def__init__(self, name):

        self.name = name

        self.params, self.grads = 

None

None

defforward(self, inputs):
raise

 NotImplementedError

defbackward(self, grad):
raise

 NotImplementedError

最基礎的一種網路層是全連線網路層,實現如下。forward 方法接收上層的輸入 inputs,實現 的運算;backward 的方法接收來自上層的梯度,計算關於引數 和輸入的梯度,然後返回關於輸入的梯度。這三個梯度的推導可以見附錄,這裡直接給出實現。w_init 和 b_init 分別是引數 和 的初始化器,這個我們在另外的一個實現初始化器中檔案 initializer.py 去實現,這部分不是核心部件,所以在這裡不展開介紹。
# layer.py
classDense(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

defforward(self, inputs):

        self.inputs = inputs

return

 inputs @ self.params[

"w"

] + self.params[

"b"

]

defbackward(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
classActivation(Layer):
"""Base activation layer"""
def__init__(self, name):

        super().__init__(name)

        self.inputs = 

None

defforward(self, inputs):

        self.inputs = inputs

return

 self.func(inputs)

defbackward(self, grad):
return

 self.derivative_func(self.inputs) * grad

deffunc(self, x):
raise

 NotImplementedError

defderivative_func(self, x):
raise

 NotImplementedError

classReLU(Activation):
"""ReLU activation function"""
def__init__(self):

        super().__init__(

"ReLU"

)

deffunc(self, x):
return

 np.maximum(x, 

0.0

)

defderivative_func(self, x):
return

 x > 

0.0

net

上文提到 net 類負責管理 tensor 在 layers 之間的前向和反向傳播。forward 方法很簡單,按順序遍歷所有層,每層計算的輸出作為下一層的輸入;backward 則逆序遍歷所有層,將每層的梯度作為下一層的輸入。這裡我們還將每個網路層引數的梯度儲存下來返回,後面引數更新需要用到。另外 net 類還實現了獲取引數、設定引數、獲取梯度的介面,也是後面引數更新時需要用到
# net.py
classNet(object):
def__init__(self, layers):

        self.layers = layers

defforward(self, inputs):
for

 layer 

in

 self.layers:

            inputs = layer.forward(inputs)

return

 inputs

defbackward(self, grad):

        all_grads = []

for

 layer 

in

 reversed(self.layers):

            grad = layer.backward(grad)

            all_grads.append(layer.grads)

return

 all_grads[::

-1

]

defget_params_and_grads(self):
for

 layer 

in

 self.layers:

yield

 layer.params, layer.grads

defget_parameters(self):
return

 [layer.params 

for

 layer 

in

 self.layers]

defset_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
classBaseLoss(object):
defloss(self, predicted, actual):
raise

 NotImplementedError

defgrad(self, predicted, actual):
raise

 NotImplementedError

classCrossEntropyLoss(BaseLoss):
defloss(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

defgrad(self, predicted, actual):

        m = predicted.shape[

0

]

        grad = np.copy(predicted)

        grad -= actual

return

 grad / m

optimizer

optimizer 主要實現一個介面 compute_step,這個方法根據當前的梯度,計算返回實際最佳化時每個引數改變的步長。我們在這裡實現常用的 Adam 最佳化器。
# optimizer.py
classBaseOptimizer(object):
def__init__(self, lr, weight_decay):

        self.lr = lr

        self.weight_decay = weight_decay

defcompute_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

classAdam(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
classModel(object):
def__init__(self, net, loss, optimizer):

        self.net = net

        self.loss = loss

        self.optimizer = optimizer

defforward(self, inputs):
return

 self.net.forward(inputs)

defbackward(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

defapply_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]

3
『整體結構』
最後我們實現出來核心程式碼部分檔案結構如下

tinynn

├── core

│ ├── initializer.py

│ ├── layer.py

│ ├── loss.py

│ ├── model.py

│ ├── net.py

│ └── optimizer.py

其中 initializer.py 這個模組上面沒有展開講,主要實現了常見的引數初始化方法(零初始化、Xavier 初始化、He 初始化等),用於給網路層初始化引數。
4
『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 實現了相同的網路結構、採用相同的採數初始化方法、最佳化器配置等等,得到的結果如下

# 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 稍微好一點點。
5
『總結』
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 等最佳化器,並且增加了動態調節學習率 LRScheduler
  • 實現了 mnist(分類)、nn_paint(迴歸)、DQN(強化學習)、AutoEncoder 和 DCGAN (無監督)等常見模型。見 tinynn/examples:https://github.com/borgwang/tinynn/tree/master/examples
tinynn 還有很多可以繼續完善的地方受限於時間還沒有完成,筆者在空閒時間會進行維護和更新。
當然 tinynn 只是一個「玩具」版本的深度學習框架,一個成熟的深度學習框架至少還需要:支援自動求導、高運算效率(靜態語言加速、支援 GPU 加速)、提供豐富的演算法實現、提供易用的介面和詳細的文件等等。這個小專案的出發點更多地是學習,在設計和實現 tinynn 的過程中筆者個人學習確實到了很多東西,包括如何抽象、如何設計元件介面、如何更效率的實現、演算法的具體細節等等。對筆者而言寫這個小框架除了瞭解深度學習框架的設計與實現之外還有一個好處:後續可以在這個框架上快速地實現一些新的演算法,新的引數初始化方法,新的最佳化演算法,新的網路結構設計,都可以快速地在這個小框架上進行實驗。如果你對自己設計實現一個深度學習框架也感興趣,希望看完這篇文章會對你有所幫助,也歡迎大家提 PR 一起貢獻程式碼~
6
『附錄: 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
技術交流群邀請函
△長按新增小助手
掃描二維碼新增小助手微信
請備註:姓名-學校/公司-研究方向
(如:小張-哈工大-對話系統)
即可申請加入自然語言處理/Pytorch等技術交流群

關於我們

MLNLP 社群是由國內外機器學習與自然語言處理學者聯合構建的民間學術社群,目前已經發展為國內外知名的機器學習與自然語言處理社群,旨在促進機器學習,自然語言處理學術界、產業界和廣大愛好者之間的進步。
社群可以為相關從業者的深造、就業及研究等方面提供開放交流平臺。歡迎大家關注和加入我們。

相關文章