動手實現一個帶自動微分的深度學習框架

轉自：Automatic Differentiation Tutorial

參考代碼：https://github.com/borgwang/tinynn-autograd (主要看 core/tensor.py 和 core/ops.py)

在這里插入圖片描述

簡介

梯度下降（Gradient Descent）及其衍生算法是神經網絡訓練的基礎，梯度下降本質上就是求解損失關于網絡參數的梯度，不斷計算這個梯度對網絡參數進行更新。現代的神經網絡框架都實現了自動求導的功能，只需要要定義好網絡前向計算的邏輯，在運算時自動求導模塊就會自動把梯度算好，不用自己手寫求導梯度。

筆者在之前的一篇文章中講解和實現了一個迷你的神經網絡框架 tinynn，在 tinynn 中我們定義了網絡層 layer 的概念，整個網絡是由一層層的 layer 疊起來的（全連接層、卷積層、激活函數層、Pooling 層等等），如下圖所示

在這里插入圖片描述

在實現的時候需要顯示為每層定義好前向 forward 和反向 backward（梯度計算）的計算邏輯。從本質上看這些 layer 其實是一組基礎算子的組合，而這些基礎算子（加減乘除、矩陣變換等等）的導函數本身都比較簡單，如果能夠將這些基礎算子的導函數寫好，同時把不同算子之間連接邏輯記錄（計算依賴圖）下來，那么這個時候就不再需要自己寫反向了，只需要計算損失，然后從損失函數開始，讓梯度自己用預先定義好的導函數，沿著計算圖反向流動即可以得到參數的梯度，這個就是自動求導的核心思想。tinynn 中之所有 layer 這個概念，一方面是符合我們直覺上的理解，另一方面是為了在沒有自動求導的情況下方便實現。有了自動求導，我們可以拋開 layer 這個概念，神經網絡的訓練可以抽象為定義好一個網絡的計算圖，然后讓數據前向流動，讓梯度自動反向流動（ TensorFlow 這個名字起得相當有水準）。

我們可以看看 PyTorch 的一小段核心的訓練代碼（來源官方文檔 MNIST 例子）

for batch_idx, (data, target) in enumerate(train_loader):data, target = data.to(device), target.to(device)optimizer.zero_grad()  # 初始化梯度output = model(data)  # 從 data 到 output 的計算圖loss = F.nll_loss(output, target) # 從 output 到 loss 的計算圖loss.backward()  # 梯度從 loss 開始反向流動optimizer.step()  # 使用梯度對參數更新

可以看到 PyTorch 的基本思路和我們上面描述的是一致的，定義好計算圖 -> forward 得到損失 -> 梯度反向流動。

自動求導設計

知道了自動求導的基本流程之后，我們考慮如何來實現。先考慮沒有自動求導，為每個運算手動寫 backward 的情況，在這種情況下我們實際上定義了兩個計算圖，一個前向一個反向，考慮最簡單的線性回歸的運算 $W X + B$ ，其計如下所示。

在這里插入圖片描述

可以看到這兩個計算圖的結構實際上是一樣的，只是在前向流動的是計算的中間結果，反向流動的是梯度，以及中間的運算反向的時候是導數運算。實際上我們可以把兩者結合到一起，只定義一次前向計算圖，讓反向計算圖自動生成

在這里插入圖片描述

從實現的角度看，如果我們不需要自動求導，那么網絡框架中的 Tensor 類只需要對 Tensor 運算符有定義，能夠進行數值運算（tinynn 中就簡單的使用 ndarray 作為 Tensor 的實現）。但如果要實現自動求導，那么 Tensor 類需要額外做幾件事：

增加一個梯度的變量保存當前 tensor 的梯度
保存當前 tensor 依賴的 tensor（如上圖中 O1O1 依賴于 X,WX,W）
保存下對各個依賴 tensor 的導函數（這個導函數的作用是將當前 tensor 的梯度傳到依賴的 tensor 上）

自動求導實現

我們按照上面的分析開始實現 Tensor 類如下，初始化方法中首先把 tensor 的值保存下來，然后有一個 requires_grad 的 bool 變量表明這個 tensor 是不是需要求梯度，還有一個 dependency 的列表用于保存該 tensor 依賴的 tensor 以及對于他們的導函數。

zero_grad() 方法比較簡單，將當前 tensor 的梯度設置為 0，防止梯度的累加。自動求導從調用計算圖的最后一個節點 tensor 的 backward() 方法開始（在神經網絡中這個節點一般是 loss）。backward() 方法主要流程為

確保改 tensor 確實需要求導 self.requires_grad == True
將從上個 tensor 傳進來的梯度加到自身梯度上，如果沒有（反向求導的起點 tensor），則將梯度初始化為 1.0
對每一個依賴的 tensor 運行保存下來的導函數，計算傳播到依賴 tensor 的梯度，然后調用依賴 tensor 的 backward() 方法。可以看到這其實就是 Depth-First Search 計算圖的節點

def as_tensor(obj):if not isinstance(obj, Tensor):obj = Tensor(obj)return objclass Tensor:def __init__(self, values, requires_grad=False, dependency=None):self._values = np.array(values)self.shape = self.values.shapeself.grad = Noneif requires_grad:self.zero_grad()self.requires_grad = requires_gradif dependency is None:dependency = []self.dependency = dependency@propertydef values(self):return self._values@values.setterdef values(self, new_values):self._values = np.array(new_values)self.grad = Nonedef zero_grad(self):self.grad = np.zeros(self.shape)def backward(self, grad=None):assert self.requires_grad, "Call backward() on a non-requires-grad tensor."grad = 1.0 if grad is None else gradgrad = np.array(grad)# accumulate gradientself.grad += grad# propagate the gradient to its dependenciesfor dep in self.dependency:grad_for_dep = dep["grad_fn"](grad)dep["tensor"].backward(grad_for_dep)

可能看到這里讀者可能會疑問，一個 tensor 依賴的 tensor 和對他們的導函數（也就是 dependency 里面的東西）從哪里來？似乎沒有哪一個方法在做保存依賴這件事。

假設我們可能會這樣使用我們的 Tensor 類

W = Tensor([[1], [3]], requires_grad=True)  # 2x1 tensor
X = Tensor([[1, 2], [3, 4], [5, 6], [7, 8]], requires_grad=True)  # 4x2 tensor
O = X @ W  # suppose to be a 4x1 tensor

如何讓 X 和 W 完成矩陣乘法輸出正確的 O 的同時，讓 O 能記下他依賴于 W 和 X 呢？答案是重載運算符。

class Tensor:# ...def __matmul__(self, other):# 1. calculate forward valuesvalues = self.values @ other.values# 2. if output tensor requires_gradrequires_grad = ts1.requires_grad or ts2.requires_grad# 3. build dependency listdependency = []if self.requires_grad:# O = X @ W# D_O / D_X = grad @ W.Tdef grad_fn1(grad):return grad @ other.values.Tdependency.append(dict(tensor=self, grad_fn=grad_fn1))if other.requires_grad:# O = X @ W# D_O / D_W = X.T @ graddef grad_fn2(grad):return self.values.T @ graddependency.append(dict(tensor=other, grad_fn=grad_fn2))return Tensor(values, requires_grad, dependency)# ...

關于 Python 中如何重載運算符這里不展開，讀者有興趣可以參考官方文檔或者這篇文章。基本上在 Tensor 類內定義了 __matmul__ 這個方法后，實際上是重載了矩陣乘法運算符 @ （Python 3.5 以上支持）。當運行 X @ W 時會自動調用 X 的 __matmul__ 方法。

這個方法里面做了三件事：

計算矩陣乘法結果（這個是必須的）
確定是否需要新生成的 tensor 是否需要梯度，這個由兩個操作數決定。比如在這個例子中，如果 W 或者 X 需要梯度，那么生成的 O也是需要計算梯度的（這樣才能夠計算 W 或者 X 的梯度）
建立 tensor 的依賴列表

自動求導中最關鍵的部分就是在這里了，還是以 O = X @ W 為例子，這里我們會先檢查是否 X需要計算梯度，如果需要，我們需要把導函數 D_O / D_X 定義好，保存下來；同樣的如果 W 需要梯度，我們將 D_O / D_W 定義好保存下來。最后生成一個 dependency 列表保存著在新生成的 tensor O 中。

然后我們再回顧前面講的 backward()方法，backward() 方法會遍歷 tensor 的 dependency ，將用保存的 grad_fn 計算要傳給依賴 tensor 的梯度，然后調用依賴 tensor 的 backward() 方法將梯度傳遞下去，從而實現了梯度在整個計算圖的流動。

grad_for_dep = dep["grad_fn"](grad)
dep["tensor"].backward(grad_for_dep)

自動求導講到這里其實已經基本沒有什么新東西，剩下的工作就是以類似的方法大量地重載各種各樣的運算符，使其能夠 cover 住大部分所需要的操作（基本上照著 NumPy 的接口都給重載一次就差不多了 🤨）。無論你定義了多復雜的運算，只要重載了相關的運算符，就都能夠自動求導了，再也不用自己寫梯度了。

在這里插入圖片描述

一個例子

大量的重載運算符的工作在文章里就不貼上來了（過程不怎么有趣），我寫在了一個 notebook 上，大家有興趣可以去看看 borgwang/toys/ml-autograd。在這個 notebook 里面重載了實現一個簡單的線性回歸需要的幾種運算符，以及一個線性回歸的例子。這里把例子和結果貼上來

# training data
x = Tensor(np.random.normal(0, 1.0, (100, 3)))
coef = Tensor(np.random.randint(0, 10, (3,)))
y = x * coef - 3params = {"w": Tensor(np.random.normal(0, 1.0, (3, 3)), requires_grad=True),"b": Tensor(np.random.normal(0, 1.0, 3), requires_grad=True)
}learng_rate = 3e-4
loss_list = []
for e in range(101):# set gradient to zerofor param in params.values():param.zero_grad()# forwardpredicted = x @ params["w"] + params["b"]err = predicted - yloss = (err * err).sum()# backward automaticallyloss.backward()# updata parameters (gradient descent)for param in params.values():param -= learng_rate * param.gradloss_list.append(loss.values)if e % 10 == 0:print("epoch-%i \tloss: %.4f" % (e, loss.values))

epoch-0 	loss: 8976.9821
epoch-10 	loss: 2747.4262
epoch-20 	loss: 871.4415
epoch-30 	loss: 284.9750
epoch-40 	loss: 95.7080
epoch-50 	loss: 32.9175
epoch-60 	loss: 11.5687
epoch-70 	loss: 4.1467
epoch-80 	loss: 1.5132
epoch-90 	loss: 0.5611
epoch-100 	loss: 0.2111

在這里插入圖片描述

接口和 PyTorch 相似，在每個循環里面首先將參數梯度設為 0 ，然后定義計算圖，然后從 loss 開始反向傳播，最后更新參數。從結果可以看到 loss 隨著訓練進行非常漂亮地下降，說明我們的自動求導按照我們的設想 work 了。

總結

本文實現了討論了自動求導的設計思路和整個過程是怎么運作的。總結起來：自動求導就是在定義了一個有狀態的計算圖，該計算圖上的節點不僅保存了節點的前向運算，還保存了反向計算所需的上下文信息。利用上下文信息，通過圖遍歷讓梯度在圖中流動，實現自動求節點梯度。

我們通過重載運算符實現了一個支持自動求導的 Tensor 類，用一個簡單的線性回歸 demo 測試了自動求導。當然這只是最基本的能實現自動求導功能的 demo，從實現的角度上看還有很多需要優化的地方（內存開銷、運算速度等），筆者有空會繼續深入研究，讀者如果有興趣也可以自行查閱相關資料。Peace out. 🤘