深度學習利器之自動微分(2)

羅西的思考發表於2021-10-14

深度學習

深度學習利器之自動微分(2)

0x00 摘要

本文和上文以 Automatic Differentiation in Machine Learning: a Survey為基礎，逐步分析自動微分這個機器學習的基礎利器。

深度學習框架，幫助我們解決的核心問題就是兩個:

反向傳播時的自動梯度計算和更新，也就是自動微分。
使用 GPU 進行計算。

本文就具體介紹自動微分究竟是什麼。

0x01 前情回顧

上文深度學習利器之自動微分(1)中，我們介紹了一些基礎概念，常見微分方法和自動微分的數學基礎。

在數學與計算代數學中，自動微分或者自動求導（Automatic Differentiation，簡稱AD）也被稱為微分演算法或數值微分。它是一種數值計算的方式，其功能是計算複雜函式（多層複合函式）在某一點處對某個的導數，梯度，Hessian矩陣值等等。

自動微分又是一種計算機程式，是深度學習框架的標配，是反向傳播演算法的泛化。開發者寫的任何模型都需要靠自動微分機制來計算梯度，從而更新模型。它的原理實際上是通過記錄運算順序，利用已經定義好的導數規則，生成一個與正常計算程式對偶的程式。

微分的幾種比較常用的方法如下：

手動求解法(Manual Differentiation) : 完全手動完成，依據鏈式法則解出梯度公式，帶入數值，得到梯度。
數值微分法(Numerical Differentiation) ：利用導數的原始定義，直接求解微分值。
符號微分法(Symbolic Differentiation) : 利用求導規則對錶達式進行自動計算，其計算結果是導函式的表示式而非具體的數值。即，先求解析解，然後轉換為程式，再通過程式計算出函式的梯度。
自動微分法(Automatic Differentiation) ：介於數值微分和符號微分之間的方法，採用類似有向圖的計算來求解微分值。

具體如下圖：

接下來我們就具體研究一下自動微分。

0x02 自動微分

2.1 分解計算

自動微分的精髓在於它發現了微分計算的本質：微分計算就是一系列有限的可微運算元的組合。

自動微分法被認為是對計算機程式進行非標準的解釋。自動微分基於一個事實，即每一個計算機程式，不論它有多麼複雜，都是在執行加減乘除這一系列基本算數運算，以及指數、對數、三角函式這類初等函式運算。於是自動微分先將符號微分法應用於最基本的運算元，比如常數，冪函式，指數函式，對數函式，三角函式等，然後代入數值，保留中間結果，最後再通過鏈式求導法則應用於整個函式。

通過將鏈式求導法則應用到這些運算上，我們能以任意精度自動地計算導數，而且最多隻比原始程式多一個常數級的運算。

我們以如下為例，這是原始公式

\[y=f(g(h(x)))=f(g(h(w_0)))=f(g(w_1))=f(w_2)=w_3 \\ \]

自動微分以鏈式法則為基礎，把公式中一些部分整理出來成為一些新變數，然後用這些新變數整體替換這個公式，於是得到：

\[\begin{align} y=f(g(h(x)))=f(g(h(w_0)))=f(g(w_1))=f(w_2)=w_3 \\ w_0=x\\ w_1=h(w_0)\\ w_2=g(w_1)\\ w_3=f(w_2)=y \end{align} \]

然後把這些新變數作為節點，依據運算邏輯把公式整理出一張有向無環圖（DAG）。即，原始函式建立計算圖，資料正向傳播，計算出中間節點xi，並記錄計算圖中的節點依賴關係。

因此，自動微分可以被認為是將一個複雜的數學運算過程分解為一系列簡單的基本運算，其中每一項基本運算都可以通過查表得出來。

2.2 計算模式

一般而言，存在兩種不同的自動微分模式：

前向自動微分（Forward Automatic Differentiation，也叫做 tangent mode AD）或者前向累積梯度（前向模式）。
後向自動微分（Reverse Automatic Differentiation，也叫做 adjoint mode AD）或者說反向累計梯度（反向模式）。

兩種自動微分模式都通過遞迴方式來求 dy/dx，只不過根據鏈式法則展開的形式不太一樣。

前向梯度累積會指定從內到外的鏈式法則遍歷路徑，即先計算 $d_{w1}/d_x$，再計算 $d_{w2}/d_{w1}$，最後計算 $dy/dw_2$，即，前向模式是在計算圖前向傳播的同時計算微分。因此前向模式的一次正向傳播就可以計算出輸出值和導數值。

\[\frac{dw_i}{dx}=\frac{dw_i}{dw_{i-1}}\frac{dw_{i-1}}{dx} \qquad with\ w_3=y \]

反向梯度累積正好相反，它會先計算$dy/dw_2$，然後計算 $d_{w2}/d_{w1}$，最後計算 $d_{w1}/d_x$。這是我們最為熟悉的反向傳播模式，它非常符合「沿模型誤差反向傳播」這一直觀思路。即，反向模式需要對計算圖進行一次正向計算，得出輸出值，再進行反向傳播。反向模式需要儲存正向傳播的中間變數值（比如$w_i$），這些中間變數數值在反向傳播時候被用來計算導數，所以反向模式的記憶體開銷要大。

\[\frac{dy}{dw_i} = \frac{dy}{dw_{i+1}}\frac{dw_{i+1}}{dw_i} \qquad with \ w_0 = x \]

讓我們再梳理一下:

前向自動微分（tangent mode AD）和後向自動微分（adjoint mode AD）分別計算了Jacobian矩陣的一列和一行。（引自 [Hascoet2013]）

前向自動微分（tangent mode AD）可以在一次程式計算中通過鏈式法則

\[\frac{∂x^k}{∂x^0_j} = \frac{∂x^k}{∂x^{k-1}} \frac{∂x^{k-1}}{∂x^0_j} \]

遞推得到Jacobian矩陣中與單個輸入有關的部分，或者說是Jacobian矩陣的一列。

後向自動微分（adjoint mode AD）利用鏈式法則

\[\frac{∂x^L_i}{∂x^k} = \frac{∂x^L_i}{∂x^{k+1}} \frac{∂x^{k+1}}{∂x^k} \]

可以僅通過一次對計算過程的遍歷得到Jacobian矩陣的一行。但它的導數鏈式法則傳遞方向和程式執行方向相反，所以需要在程式計算過程中記錄一些額外的資訊來輔助求導，這些輔助資訊包括計算圖和計算過程的中間變數。

2.3 樣例

我們以公式 $ f(x_1, x_2) = ln(x_1) + x_1x_2 - sin(x_2)$ 為例，首先把它轉換成一個計算圖，則如下:

這裡對於變數我們使用三段法表示：

輸入變數：自變數維度為 n，這裡 n = 2，輸入變數就是 $x_1, x_2$ 。
中間變數：中間變數這裡是 $v_{-1} $ 到 $ v_5$，在計算過程中，只需要針對這些中間變數做處理即可：將符號微分法應用於最基本的運算元，然後代入數值，保留中間結果，最後再應用於整個函式。
輸出變數：假設輸出變數維度為 m，這裡 m = 1，輸出變數就是 $y_1$，也就是$f(x_1, x_2)$。

或者加上符號我們更容易理解。

轉化成如上DAG（有向無環圖）結構之後，我們可以很容易分步計算函式的值，並求取它每一步的導數值，然後，我們把 $df / dx_1$ 求導過程利用鏈式法則表示成如下的形式：

可以看出整個求導可以被拆成一系列微分運算元的組合。這裡的計算可以分為兩種：

將這個式子從前往後算的叫前向模式（Forward Mode）。
將這個式子從後往前算的叫逆向模式（Reverse Mode）。

2.4 前向模式（Forward Mode）

前向模式從計算圖的起點開始，沿著計算圖邊的方向依次向前計算，最終到達計算圖的終點。它根據自變數的值計算出計算圖中每個節點的值以及其導數值，並保留中間結果。一直得到整個函式的值和其導數值。整個過程對應於一元複合函式求導時從最內層逐步向外層求導。

2.4.1 計算過程

下面是前向模式的計算過程，下表中，左半部分是從左往右每個圖節點的求值結果和計算過程，右半部分是每個節點對 $x_1$的求導結果和計算過程。這裡 $\dot V_i$ 表示 $V_i$ 對 $x_1$的偏導數。

我們給出一個推導分析，對於節點數值的計算如下：

我們給輸入節點賦值：$v_{-1} = x_1 = 2$，$v_{0} = x_2 = 5$。
計算$v_1$節點：$v_1 = ln\ v_{-1} = ln\ x_1 = ln\ 2$ 。
計算 $v_2$ 節點，節點 v_2 依賴於$v_{-1}$ 和 $v_0$，即: $v_2=v_{−1} \times v_0=x_1 \times x_2= 2 \times 5 = 10$
計算 v3 節點，$ v_3 = sin \ v_0 = sin \ 5$
計算 $ v_4 $節點，$v_4 = v_1 + v_2 = 0.693 + 10$
計算 $v_5$節點，$v_5 = v_1 + v_2 = 10.693 + 0.959$
最終 $ y = v_5 = 11.652$

此時，我們得到了圖中所有節點的數值。而且在計算節點數值的同時，我們也一起計算導數，假設我們求 $\frac{∂y}{∂x_{1}}$我們也是從輸入計算。

$v_{-1}$ 節點對於$x_1$梯度 : 因為 $v_{-1} = x_1 $，所以 $\frac{∂v_{-1}}{∂x{1}} = 1$
$v_{0}$ 節點對於$x_1$梯度 : 因為 $v_{0} = x_2 $，這樣就和 $x_{1}$ 無關，所以 $\frac{∂v_{0}}{∂x{1}} = 0$
$v_{1}$ 節點對於$x_1$梯度 : $\frac{∂v_{1}}{∂x_{1}} = \frac{∂\ log \ x_1}{∂x_1} = \frac{1}{x_1} = \frac{1}{2}$
$v_{2}$ 節點對於$x_1$梯度 : $\frac{∂v_{2}}{∂x_{1}} = \frac{∂v_{-1} }{∂x_1}\times v_0 + \frac{∂v_{0}}{ ∂x_1}\times v_{-1} = 1 \times 5 + 0 \times 2 $
$v_{3}$ 節點對於$x_1$梯度 : $\frac{∂v_{0}}{∂x_{1}} \times cos \ v_0 = 0 \times cos \ 5$
$v_{4}$ 節點對於$x_1$梯度 : $\frac{∂v_{1}}{∂x_{1}} = \frac{∂ v_1}{∂x_1} + \frac{∂ v_2}{∂x_1} = 0.5 + 5$
$v_{5}$ 節點對於$x_1$梯度 : $\frac{∂v_{1}}{∂x_{1}} = \frac{∂v_4}{∂x_1} - \frac{∂v_3}{∂x_1} = 5.5 - 0$
所以 $\frac{∂y}{∂x_{1}} = \frac{∂v_5}{∂x_1} = 5.5$

請注意:

計算圖各個節點的數值和該節點的導數可同步求出。因為我們已經瞭解了深度學習，所以從深度學習角度來看，前向模式只需要完成計算出計算圖中各節點的值。求導數則可以由反向模式來實現。
注意到每一步的求導都會利用到上一步的求導結果，這樣不至於重複計算，因此也不會產生像符號微分法的"expression swell"問題。自動微分的前向模式實際上與我們在微積分裡所學的求導過程一致。

3.4.2 推廣

這個計算過程可以推廣到雅克比矩陣，假設一個函式有 n 個輸入變數 $x_i$，m個輸入變數 $y_j$，即輸入向量 $x∈R^n$，而輸出向量 $y∈R^m$，則這個函式就是對映

\[f : R ^n \rightarrow R^m \]

在這種情況下，每個自動微分的前向傳播計算時候，初始輸入被設定為 $\dot{x} = 1$，其餘被設定為 0。

對於雅克比矩陣

可以看出，一次前向計算，可以求出Jacobian矩陣的一列資料。比如 $\dot{x_3} = 1$，對應就可以求出來第3列。

即，tangent mode AD可以在一次程式計算中通過鏈式法則遞推得到Jacobian矩陣中與單個輸入有關的部分，或者說是Jacobian矩陣的一列。如下圖：

但是如果想求出來對所有輸入 $x_i\ $，$i = 1,..,n$ ，我們需要計算 n 次才能求出所有列。

進一步，通過設定 $\dot{x} = r$，我們可以在一次前向傳播中直接計算 Jacobian–vector 乘積。

3.4.3 問題

前向模式的優點在於：實現起來很簡單 [Revels2016]，也不需要很多額外的記憶體空間。

前向模式的問題在於：

每次前向計算只能計算對一個自變數的偏導數，對於一元函式求導是高效的。但是機器學習模型的引數通常有 $10^6$數量級之多。
如果有一個函式，其輸入有 n 個，輸出有 m個，對於每個輸入來說，前向模式都需要遍歷計算過程以得到這個輸入的導數，求解整個函式梯度需要 n 遍如上計算過程。尤其是神經網路這種 $n \gg m$ 的模型，前向模式就太低效了，重複遍歷計算過程 $10^6$次顯然是無法接受的。

於是86年Hinton提出了用後向傳播技術訓練神經網路 [Rumelhart1986]，也就是接下來要說的反向模式（adjoint mode AD）。

3.5 反向模式（Reverse Mode）

3.5.1 思路

後向自動微分是基於鏈式法則來工作。我們僅需要一個前向過程和反向過程就可以計算所有引數的導數或者梯度。因為需要結合前向和後向兩個過程，因此反向自動微分會將所有的操作以一張圖的方式儲存下來，這張圖就是計算圖，計算圖是一個有向無環圖（DAG），它表達了函式和變數的關係。這是深度學習框架的核心之一：如何幹淨地產生一個計算圖，隨後再高效地計算它。

反向模式根據計算圖從後（最後一個節點）向前計算，依次得到對每個中間變數節點的偏導數，直到到達自變數節點處，這樣就得到了每個輸入的偏導數。在每個節點處，根據該節點的後續節點（前向傳播中的後續節點）計算其導數值。整個過程對應於多元複合函式求導時從最外層逐步向內側求導。這樣可以有效地把各個節點的梯度計算解耦開，每次只需要關注計算圖中當前節點的梯度計算。

神經網路的backprop演算法就是reverse mode自動微分的一種特殊形式。

從下圖可以看出來，reverse mode和forward mode是一對相反過程，reverse mode從最終結果開始求導，利用最終輸出對每一個節點進行求導。下圖虛線就是反向模式。

3.5.2 計算過程

前向和後向兩種模式的過程表達如下，表的左列淺色為前向計算函式值的過程，與前向計算時相同。右面列深色為反向計算導數值的過程。

需要注意的是：左列先計算，順序是由上自下，右列後計算，順序是由下往上。右圖必須從下往上看，即，先計算輸出 y 對節點 $v_5$ 的導數，用 $\tilde{v_5}$表示 $\frac{∂y}{∂v_5}$，這樣的記號可以強調我們對當前計算結果進行快取，以便用於後續計算，而不必重複計算。

因為前向計算函式值得順序與前向模式相同，所以我們不討論，專注看看反向計算導數值的過程。為了更好的說明，我們在原圖上插入數值和公式。

反向計算導數值時，順序如下：

計算 y 對 $v_5$的導數值，因為 $ y = v_5$, 所以 $\frac{∂y}{∂v_5} = v'_5 = 1$。
計算 y 對 $v_4$的導數值，因為 $v_4$ 在圖上只有一個後續節點 $v_5$, 並且 $v_5=v_3−v_4$，所以依據鏈式法則得到$\frac{∂y}{∂v_4} =\frac{∂y}{∂v5} \times \frac{∂v5}{∂v4} = v'_5 \frac{∂v_5}{∂v_4} = v'_5 \frac{∂(v_3 - v_4)}{∂v_4} = v'_5 \times 1 = 1$，將結果寫在$v_4$指向$v_5$的邊上。
計算 y 對 $v_3$的導數值，因為 $v_3$ 在圖上只有一個後續節點 $v_5$, 並且 $v_5=v_3−v_4$，所以依據鏈式法則得到$\frac{∂y}{∂v_3} =\frac{∂y}{∂v5} \times \frac{∂v5}{∂v3} = v'_5 \frac{∂v_5}{∂v_3} = v'_5 \frac{∂(v_3 - v_4)}{∂v_3} = v'_5 \times (-1) = -1$，，將結果寫在$v_3$指向$v_5$的邊上。
計算 y 對 $v_1$的導數值，因為 $v_1$ 在圖上只有一個後續節點 $v_4$, 並且 $v_4=v_1+v_2$，所以依據鏈式法則得到$\frac{∂y}{∂v_1} =\frac{∂y}{∂v_1} \times \frac{∂v_4}{∂v_1} = v'_4 \frac{∂v_4}{∂v_1} = v'_4 \frac{∂(v_1 + v_1)}{∂v_1} = v'_4 \times 1 = 1$，，將結果寫在$v_1$指向$v_4$的邊上。
計算 y 對 $v_2$的導數值，因為 $v_2$ 在圖上只有一個後續節點 $v_4$, 並且 $v_4=v_1+v_2$，所以依據鏈式法則得到$\frac{∂y}{∂v_2} =\frac{∂y}{∂v_4} \times \frac{∂v_4}{∂v_2} = v'_4 \frac{∂v_4}{∂v_2} = v'_4 \frac{∂(v_1 + v_2)}{∂v_2} = v'_4 \times 1 = 1$，，將結果寫在$v_2$指向$v_4$的邊上。
接下來要計算 y 對 $v_0$的導數值和 y 對 $v_{-1}$的導數值，因為 $v_0$ 和$v_{-1}$ 都是後續有兩個節點，
1. $v_0$ 在圖上有兩個後續節點 $v_2$ 和 $v_3$, 並且 $v_2=v_{-1} \times v_0$ 和 $v_3= sin \ v_0$
2. $v_{-1}$ 在圖上有兩個後續節點 $v_2$ 和 $v_1$, 並且 $v_2=v_{-1} \times v_0$ 和 $v_1= ln \ v_{-1}$
所以我們只能分開計算。
計算 $\frac{∂v_3}{∂v_0} = \frac{∂sin\ v_0}{∂v_0} = cos\ v_0 = 0.284$，將結果寫在$v_0$指向$v_3$的邊上。
計算 $\frac{∂v_2}{∂v_0} = \frac{∂(v_{-1}v_0)}{∂v_0} = v_{-1} = 2$，將結果寫在$v_0$指向$v_2$的邊上。
計算 $\frac{∂v_2}{∂v_{-1}} = \frac{∂(v_{-1}v_0)}{∂v_{-1}} = v_0 = 5$，將結果寫在$v_{-1}$指向$v_2$的邊上。
計算 $\frac{∂v_1}{∂v_{-1}} = \frac{∂(ln \ v_{-1})}{∂v_{-1}} = \frac{1}{x_1} = 0.5$，將結果寫在$v_{-1}$指向$v_1$的邊上。

到目前為止，我們已經計算出來了所有步驟的偏導數的數值。現在需要計算 $\frac{∂y}{∂x_1}$和 $\frac{∂y}{∂x_2}$。計算 $\frac{∂y}{∂x_1}$ 就是從 y 開始從後向前走，走回到 $x_1$，因為有多條路徑，所以對於每一條路徑，需要將這個路徑上的數值連乘起來得到一個乘積數值，然後將這多條路徑的乘積數值相加起來，就得到了 $\frac{∂y}{∂x_1}$ 的數值。$\frac{∂y}{∂x_2}$的計算方法與此相同。

從 y 到 $x_1$ 的路徑有兩條，分別是：

$v_5→v_4→v_1→v_{−1}$ ：其數值乘積是 1∗1∗0.5=0.5。
$v_5→v_4→v_2→v_{−1}$ ：其數值乘積是 1∗1∗ 5= 5。

因此，$\frac{∂y}{∂x_1} = 0.5 + 5 = 5.5 $

從 y 到 $x_2$ 的路徑有兩條，分別是：

$v_5→v_4→v_2→v_0$ ：其數值乘積是 1∗1∗2=2.0。
$v_5→v_3→v_0$ ：其數值乘積是 (−1)∗0.284=−0.284。

因此，$\frac{∂y}{∂x_1} = 2.0+(−0.284)=1.716 $

3.5.3 推廣

如果要同時計算多個變數的偏導數，則可以藉助雅克比矩陣完成。假設一個函式有 n 個輸入變數 $x_i$，m個輸入變數 $y_j$，即輸入向量 $x∈R^n$，而輸出向量 $y∈R^m$，則這個函式就是對映

\[f : R ^n \rightarrow R^m \]

得到雅克比矩陣是：

即通過雅克比矩陣轉置與後續節點梯度值的乘積，可以得到當前節點的梯度值。

3.5.4 記憶體問題和inplace 操作

前向模式在計算之中，計算圖各個節點的數值和該節點的導數可同步求出，但是代價就是對於多個輸入需要多次計算才行。而反向模式可以通過一次反向傳輸，就計算出所有偏導數，而且中間的偏導數計算只需計算一次，減少了重複計算的工作量，在多引數的時候後向自動微分的時間複雜度更低，但這是以增加儲存量需求為代價的。

因為在反向計算時需要尋找它所有的後續節點，收集這些節點的導數值，然後計算本節點的導數值。整個計算過程中不僅利用了每個節點的後續節點的導數值$\frac{∂f}{∂v_{n_j}}$ ，還需要利用某些節點的函式值以計算 $\frac{∂v_{n_j}}{∂v_i}$，因此需要在前向計算時儲存所有節點的值，供反向計算使用，不必重複計算。

因為反向模式的特點帶來了大量記憶體佔用，為了減少記憶體操作，各個深度學習框架做了很多優化，也帶來了很多限制。比如 PyTorch 的求導就不支援絕大部分 inplace 操作。

inplace 指的是在不更改變數的記憶體地址的情況下，直接修改變數的值。

# 情景 1，不是 inplace，類似 Python 中的 `i=i+1`
a = a.exp()
# 情景 2，是 inplace 操作，類似 `i+=1`
a[0] = 10

為什麼不支援，因為如果變數 A 已經參與了正向傳播計算，然後它的數值被修改了。但接下來反向傳播時候怎麼辦？反向傳播用這個新 A 值肯定是不正確的。但是從哪裡去找修改之前的值呢？

一個辦法是對於每個變數（因為無法確定哪個變數可能被修改）在做前向傳播計算之後，都開闢新的空間來儲存這個變數數值，這樣以後無論怎麼修改都不怕了。但是這樣會導致記憶體劇增。所以只能限制 inplace 操作。

3.6 比較

如下圖所示，前向自動微分（tangent mode AD和後向自動微分（adjoint mode AD）分別計算了Jacobian矩陣的一列和一行。（引自 [Hascoet2013]）

從矩陣乘法次數的角度來看，前向模式和反向模式的不同之處在於矩陣相乘的起始之處不同。

正向，反向兩種模式算的梯度值是一樣的，但因為計算順序不同，所以計算速度有差異。

當輸出維度小於輸入維度，反向模式的乘法次數要小於前向模式。因此，當輸出的維度大於輸入的時候，適宜使用前向模式微分；當輸出維度遠遠小於輸入的時候，適宜使用反向模式微分。即，後向自動微分更加適合多引數的情況，多引數的時候後向自動微分的時間複雜度更低，只需要一遍reverse mode的計算過程，便可以求出輸出對於各個輸入的導數，從而輕鬆求取梯度用於後續優化更新。

0x04 PyTorch

我們接下來大體看看 PyTorch 的自動微分，為後續分析打下基礎。

4.1 原理

PyTorch 反向傳播的計算主要是通過autograd類實現了自動微分功能，而autograd 的基礎是：

數學基礎：複合函式，鏈式求導法則和雅克比矩陣；
工程基礎：Tensor 構成的計算圖（DAG有向無環圖）；

4.1.1 雅克比矩陣

在數學上，如果你有一個向量值函式 $\vec{y}=f(\vec{x})$ ，那麼梯度 $\vec{y}$ 關於 $\vec{x}$ 的梯度是一個雅可比矩陣 J，雅可比矩陣 J 包含以下所有偏導組合：

\[J = [\frac{∂f}{∂x_1} ... \frac{∂f}{∂x_n}] = \left\{ \begin{matrix} \frac{∂f_1}{∂x_1} & \cdots & \frac{∂f_1}{∂x_n}\\ \vdots & \ddots & \vdots \\ \frac{∂f_m}{∂x_1} & \cdots & \frac{∂f_m}{∂x_n} \end{matrix} \right\} \]

或者是：

\[J = \left(\begin{array}{cc} \frac{\partial \bf{y}}{\partial x_{1}} & ... & \frac{\partial \bf{y}}{\partial x_{n}} \end{array}\right) = \left(\begin{array}{ccc} \frac{\partial y_{1}}{\partial x_{1}} & \cdots & \frac{\partial y_{1}}{\partial x_{n}}\\ \vdots & \ddots & \vdots\\ \frac{\partial y_{m}}{\partial x_{1}} & \cdots & \frac{\partial y_{m}}{\partial x_{n}} \end{array}\right) \]

上面的矩陣表示f(X)相對於X的梯度。注意：雅可比矩陣實現的是 n 維向量到 m 維向量的對映。

我們下面看看 PyTorch 的思路。

backward 函式

在現實中，PyTorch 是使用backward函式進行反向求導。

def backward(self, gradient=None,...)

backward 方法根據gradient引數的不同有2種可能：

如果gradient 是標量，我們可以直接計算完整的雅克比矩陣。如果預設則使用預設值 1，即 backward(torch.tensor(1.))。
如果 gradient 是向量 vector。PyTorch 不會顯式構造整個雅可比矩陣（實際的梯度），而是直接計算 Jacobian 乘積（Jacobian vector product），這通常更簡單有效，而且儲存真正的梯度會浪費大量空間。

torch.autograd 類

backward 方法最終呼叫的是torch.autograd 類。在 PyTorch 之中，torch.autograd 類從數學來說就是一個雅可比向量積計算引擎。

我們假設如下：一個啟用梯度的張量X，$X = [x_1,x_2,…,x_n]$, 假設這是某個機器學習模型的權值。X 經過一些運算形成一個向量 Y ，$Y = f(X) = [y_1, y_2，…,y_m]$。然後使用Y計算標量損失l。假設向量 v 恰好是標量損失 l 關於向量 Y 的梯度，則向量 v 稱為grad_tensor（梯度張量），

對於一個向量輸入$\vec{v}$，backward方法計算的是 $ J^{T}\cdot \vec{v}$ 。引數 grad_tensor 就是這裡的 v，需要與 Tensor 本身有相同的size。如果 $\vec{v}$ 恰好是標量函式的梯度 $l=g\left(\vec{y}\right)$，即當gradient為標量時候，

\[\vec{v} = \left(\begin{array}{ccc}\frac{\partial l}{\partial y_{1}} & \cdots & \frac{\partial l}{\partial y_{m}}\end{array}\right)^{T} \]

損失函式 $l=g\left(\vec{y}\right)$ 是一個從向量 $\vec{y}$ 到標量 l 的對映，那麼 l 對於 y 的梯度是 $(\frac{∂l}{∂y_1} ... \frac{∂l}{∂y_m})$。根據鏈式法則，$l = g(\vec{y})$ 和 $\vec{y} = f(\vec{x})$ 則標量 l 關於 $\vec{x}$的梯度就是向量-雅可比積：

\[J^{T}\cdot \vec{v}=\left(\begin{array}{ccc} \frac{\partial y_{1}}{\partial x_{1}} & \cdots & \frac{\partial y_{m}}{\partial x_{1}}\\ \vdots & \ddots & \vdots\\ \frac{\partial y_{1}}{\partial x_{n}} & \cdots & \frac{\partial y_{m}}{\partial x_{n}} \end{array}\right)\left(\begin{array}{c} \frac{\partial l}{\partial y_{1}}\\ \vdots\\ \frac{\partial l}{\partial y_{m}} \end{array}\right)=\left(\begin{array}{c} \frac{\partial l}{\partial x_{1}}\\ \vdots\\ \frac{\partial l}{\partial x_{n}} \end{array}\right) \]

這種計算雅可比矩陣並將其與向量v相乘的方法使PyTorch能夠輕鬆地為非標量輸出提供外部梯度。

4.1.2 計算圖

前面提到過，訓練分幾步執行：

原始函式建立計算圖，將問題轉化成一種有向無環圖。
進行正向傳播，計算出中間節點，並記錄計算圖中的節點依賴關係。
反向傳播，從輸出開始遍歷計算圖，計算輸出對於每個節點的導數。

可以看出來，計算圖是這裡的關鍵，是自動微分的工程基礎。計算圖就是用圖的方式來表示計算過程，每一個資料都是計算圖的一個節點，資料之間的計算即流向關係由節點之間的邊來表示。它將多輸入的複雜計算表達成了由多個基本二元計算組成的有向圖，並保留了所有中間變數，這種結構天然適用於利用鏈式法則進行自動求導，可以完全向使用者隱藏求導過程。

比如該圖所表示的運算為 $(a+b) \times c$。其中節點 $v_1$ 表示中間結果，$v_2$ 表示最終結果。計算圖的一個核心特徵是可以通過傳遞"區域性計算”來獲得最終結果。

深度學習框架中，底層結構都是由張量組成的計算圖，當然PyTorch在實際前向傳播過程中，並沒有顯示地構造出計算圖，但是其計算路徑的確是沿著計算圖的路徑來進行，而向後圖是由autograd類在向前傳遞過程中自動動態建立的。

4.1.3 神經網路中的鏈式法則

下面我們看一個簡單的神經網路模型中鏈式求導法則應用的例子，摘錄自 https://blog.paperspace.com/pytorch-101-understanding-graphs-and-automatic-differentiation/：

某神經網路中有5個神經元，其中 $w_1$至 $w_4$ 是權重矩陣，L 是輸出，其計算關係如下：

\[b=w_1∗a \\ c=w_2∗a\\ d=w_3∗b+w_4∗c\\ L=10−d \]

則其組成的前向計算圖如下：

其求導規則是：

\[\frac{∂L}{∂w_4} = \frac{\partial{L}}{\partial{d}} * \frac{\partial{d}}{\partial{w_4}}\\ \frac{\partial{L}}{\partial{w_3}} = \frac{\partial{L}}{\partial{d}} * \frac{\partial{d}}{\partial{w_3}} \\ \frac{\partial{L}}{\partial{w_2}} = \frac{\partial{L}}{\partial{d}} * \frac{\partial{d}}{\partial{c}} * \frac{\partial{c}}{\partial{w_2}} \\ \frac{\partial{L}}{\partial{w_1}} = \frac{\partial{L}}{\partial{d}} * \frac{\partial{d}}{\partial{b}} * \frac{\partial{b}}{\partial{w_1}} \]

對應如下圖：

在pytorch中，反向傳播計算過程就和上圖類似。

比如，如果計算 L 對 w2 的偏導，就是按照如下計算圖的路徑：

\[\frac{\partial{L}}{\partial{w_2}} = \frac{\partial{L}}{\partial{d}} * \frac{\partial{d}}{\partial{c}} * \frac{\partial{c}}{\partial{w_2}} \\ \]

即，先求 L 對於 d 的偏導數，再求 d 對 c 的偏導數，再求 c 對 $w_2$ 的偏導數，最後把所有乘起來，就是最終答案。

另外，雅克比矩陣可表示所有L對所有權重的偏導：

\[J = [\frac{\partial{L}}{\partial{w_1}}, \frac{\partial{L}}{\partial{w_2}}, \frac{\partial{L}}{\partial{w_3}}, \frac{\partial{L}}{\partial{w_4}}] \]

反向傳播過程中一般用來傳遞上游傳來的梯度，從而實現鏈式法則。

I可以表示損失函式，y表示輸出，x表示中間的資料或者權重引數。中間層可以類比，每一層都可以看做是一個雅克比矩陣，v可能就是一個矩陣了。通過迭代可以求得loss對任意引數的導數。

雅克比矩陣的這個特點使得將外部梯度輸入到一個帶有非標量輸出的模型變得非常簡單。

4.2 PyTorch 功能

PyTorch 有兩種求導方法。

4.2.1 方法1 backward

backward 是通過 torch.autograd.backward(）求導。具體定義如下：

    def backward(self, gradient=None, retain_graph=None, create_graph=False, inputs=None):
        r"""
        Args:
            gradient (Tensor or None): Gradient w.r.t. the
                tensor. If it is a tensor, it will be automatically converted
                to a Tensor that does not require grad unless ``create_graph`` is True.
                None values can be specified for scalar Tensors or ones that
                don't require grad. If a None value would be acceptable then
                this argument is optional.
            retain_graph (bool, optional): If ``False``, the graph used to compute
                the grads will be freed. Note that in nearly all cases setting
                this option to True is not needed and often can be worked around
                in a much more efficient way. Defaults to the value of
                ``create_graph``.
            create_graph (bool, optional): If ``True``, graph of the derivative will
                be constructed, allowing to compute higher order derivative
                products. Defaults to ``False``.
            inputs (sequence of Tensor): Inputs w.r.t. which the gradient will be
                accumulated into ``.grad``. All other Tensors will be ignored. If not
                provided, the gradient is accumulated into all the leaf Tensors that were
                used to compute the attr::tensors. All the provided inputs must be leaf
                Tensors.
        """
        if has_torch_function_unary(self):
            return handle_torch_function(
                Tensor.backward,
                (self,),
                self,
                gradient=gradient,
                retain_graph=retain_graph,
                create_graph=create_graph,
                inputs=inputs)
        torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)

其引數為：

gradient: 如果張量是非標量（即其資料有多個元素）且需要梯度，則backward函式還需要指定“梯度”。它應該是型別和位置都匹配的張量。
retain_graph: 如果設定為False，用於計算梯度的圖形將被釋放。通常在呼叫backward後，會自動把計算圖銷燬，如果要想對某個變數重複呼叫backward，需要將該引數設定為True。
create_graph: 當設定為True的時候可以用來計算更高階的梯度。
inputs ：需要計算梯度的張量。如果沒有提供，則梯度被累積到所有葉子張量上。

需要注意的是，這個函式只是提供求導功能，並不返回值，返回的總是None。

簡單的自動求導

如果Tensor類表示的是一個標量（即它包含一個元素的張量），則不需要為backward()指定任何引數。

a = torch.tensor(1.0, requires_grad=True)
b = torch.tensor(2.0, requires_grad=True)
z = x**3+b
z.backward()
print(z, a.grad, b.grad)

輸出

tensor([ 8., 10., 12.])
tensor([147., 192., 243.])

如果使用inputs引數，比如：

torch.autograd.backward([z], inputs=[a])

則只會在 a 上累積，b上不會計算梯度。

複雜的自動求導

但是如果需要計算梯度的Tensor類有更多的元素，即張量是一個向量或者是一個矩陣，則需要指定一個gradient引數，它是形狀匹配的張量，比如：

import torch

x = torch.randn(2, 2, dtype=torch.double, requires_grad=True)
y = torch.randn(2, 2, dtype=torch.double, requires_grad=True)

def fn():
    return x ** 2 + y * x + y ** 2

gradient = torch.ones(2, 2)

torch.autograd.backward(fn(), gradient, inputs=[x, y])

print(x.grad)
print(y.grad)

輸出

tensor([[-1.0397, -2.4018],
        [-0.5114,  0.2455]], dtype=torch.float64)
tensor([[-0.9240, -2.5764],
        [-1.4938,  1.2254]], dtype=torch.float64)

4.2.2 方法2 grad

方法二是通過torch.autograd.grad(）來求導。該函式會自動完成求導過程，而且會自動返回對於每一個自變數求導的結果。這是和backward不一樣的地方。

具體定義如下：

def grad(
    outputs: _TensorOrTensors,
    inputs: _TensorOrTensors,
    grad_outputs: Optional[_TensorOrTensors] = None,
    retain_graph: Optional[bool] = None,
    create_graph: bool = False,
    only_inputs: bool = True,
    allow_unused: bool = False
) -> Tuple[torch.Tensor, ...]:
    r"""Computes and returns the sum of gradients of outputs with respect to
    the inputs.

    ``grad_outputs`` should be a sequence of length matching ``output``
    containing the "vector" in Jacobian-vector product, usually the pre-computed
    gradients w.r.t. each of the outputs. If an output doesn't require_grad,
    then the gradient can be ``None``).

    If ``only_inputs`` is ``True``, the function will only return a list of gradients
    w.r.t the specified inputs. If it's ``False``, then gradient w.r.t. all remaining
    leaves will still be computed, and will be accumulated into their ``.grad``
    attribute.

    .. note::

        If you run any forward ops, create ``grad_outputs``, and/or call ``grad``
        in a user-specified CUDA stream context, see
        :ref:`Stream semantics of backward passes<bwd-cuda-stream-semantics>`.

    Args:
        outputs (sequence of Tensor): outputs of the differentiated function.
        inputs (sequence of Tensor): Inputs w.r.t. which the gradient will be
            returned (and not accumulated into ``.grad``).
        grad_outputs (sequence of Tensor): The "vector" in the Jacobian-vector product.
            Usually gradients w.r.t. each output. None values can be specified for scalar
            Tensors or ones that don't require grad. If a None value would be acceptable
            for all grad_tensors, then this argument is optional. Default: None.
        retain_graph (bool, optional): If ``False``, the graph used to compute the grad
            will be freed. Note that in nearly all cases setting this option to ``True``
            is not needed and often can be worked around in a much more efficient
            way. Defaults to the value of ``create_graph``.
        create_graph (bool, optional): If ``True``, graph of the derivative will
            be constructed, allowing to compute higher order derivative products.
            Default: ``False``.
        allow_unused (bool, optional): If ``False``, specifying inputs that were not
            used when computing outputs (and therefore their grad is always zero)
            is an error. Defaults to ``False``.
    """
    outputs = (outputs,) if isinstance(outputs, torch.Tensor) else tuple(outputs)
    inputs = (inputs,) if isinstance(inputs, torch.Tensor) else tuple(inputs)
    overridable_args = outputs + inputs
    if has_torch_function(overridable_args):
        return handle_torch_function(
            grad,
            overridable_args,
            outputs,
            inputs,
            grad_outputs=grad_outputs,
            retain_graph=retain_graph,
            create_graph=create_graph,
            only_inputs=only_inputs,
            allow_unused=allow_unused,
        )

    if not only_inputs:
        warnings.warn("only_inputs argument is deprecated and is ignored now "
                      "(defaults to True). To accumulate gradient for other "
                      "parts of the graph, please use torch.autograd.backward.")

    grad_outputs_ = _tensor_or_tensors_to_tuple(grad_outputs, len(outputs))
    grad_outputs_ = _make_grads(outputs, grad_outputs_)

    if retain_graph is None:
        retain_graph = create_graph

    return Variable._execution_engine.run_backward(
        outputs, grad_outputs_, retain_graph, create_graph,
        inputs, allow_unused, accumulate_grad=False)

引數如下：

outputs: 結果節點，微分函式的輸出，即需要求導的那個函式。
inputs: 葉子節點，即返回的梯度，即函式的自變數。
grad_outputs: Jacobian-vector 積中的向量。
retain_graph: 如果設定為False，用於計算梯度的圖形將被釋放。通常在呼叫backward後，會自動把計算圖銷燬，如果要想對某個變數重複呼叫backward，需要將該引數設定為True。
create_graph: 當設定為True的時候可以用來計算更高階的梯度。
allow_unused: 預設為False, 即必須要指定input,如果沒有指定的話則報錯。

例子如下：

import torch

x = torch.randn(2, 2, requires_grad=True)
y = torch.randn(2, 2, requires_grad=True)
z = x ** 2 + y * x + y ** 2
z.backward(torch.ones(2, 2), create_graph=True)

x_grad = 2 * x + y
y_grad = x + 2 * y

grad_sum = 2 * x.grad + y.grad
x_hv = torch.autograd.grad(
    outputs=[grad_sum], grad_outputs=[torch.ones(2, 2)],
    inputs=[x], create_graph=True)

print(x.grad)
print(y.grad)

輸出

tensor([[ 2.3553, -1.9640],
        [ 0.5467, -3.1051]], grad_fn=<CopyBackwards>)
tensor([[ 1.8319, -2.8185],
        [ 0.5835, -2.5158]], grad_fn=<CopyBackwards>)

4.3 模擬印證

接下來，我們用 PyTorch 程式碼來大致模擬印證一下自動微分的過程。

4.3.1 原始版本

以下是原始版本，就是直接用一個公式來計算出梯度。

import torch
x1 = torch.tensor(2., requires_grad=True)
x2 = torch.tensor(5., requires_grad=True)

y = torch.log(x1) + x1 * x2 - torch.sin(x2)
grads = torch.autograd.grad(y, [x1, x2])

print('y is :', y)
print('grad is : ', grads[0],grads[1]) #

輸出為如下，其中 grads[0],grads[1] 分別是y 對於 x1 和 x2 的梯度。

y is : tensor(11.6521, grad_fn=<SubBackward0>)
grad is :  tensor(5.5000) tensor(1.7163)

4.3.2 前向模式

下面我們看看前向模式的大致模擬，我們把中間計算用變數表示出來，其中 dv1~dv4 就是梯度。

可以看到在前向過程中的變數數值和梯度數值，大家可以和前面推導的計算過程印證看看。

import torch
x1 = torch.tensor(2., requires_grad=True)
x2 = torch.tensor(5., requires_grad=True)

v1 = torch.log(x1)
v2 = x1 * x2
v3 = torch.sin(x2)
v4 = v1 + v2
y = v4 - v3

dv1 = torch.autograd.grad(v1, [x1], retain_graph=True)
dv2 = torch.autograd.grad(v2, [x1], retain_graph=True)
dv3 = torch.autograd.grad(v3, [x2], retain_graph=True)
dv4 = torch.autograd.grad(v4, [x1, x2], retain_graph=True)

grads = torch.autograd.grad(y, [x1, x2])

print('v1 is :', v1)
print('v2 is :', v2)
print('v3 is :', v3)
print('v4 is :', v4)
print('y is :', y)

print('dv1 is :', dv1)
print('dv2 is :', dv2)
print('dv3 is :', dv3)
print('dv4 is :', dv4)
print('grad is : ', grads[0],grads[1])

輸出是：

v1 is : tensor(0.6931, grad_fn=<LogBackward>)
v2 is : tensor(10., grad_fn=<MulBackward0>)
v3 is : tensor(-0.9589, grad_fn=<SinBackward>)
v4 is : tensor(10.6931, grad_fn=<AddBackward0>)
y is : tensor(11.6521, grad_fn=<SubBackward0>)
dv1 is : (tensor(0.5000),)
dv2 is : (tensor(5.),)
dv3 is : (tensor(0.2837),)
dv4 is : (tensor(5.5000), tensor(2.))
grad is :  tensor(5.5000) tensor(1.7163)

4.3.3 反向模式

下面我們看看反向模式的大致模擬，我們把中間計算用變數表示出來，其中 dv系列變數就是梯度。

可以看到在反向過程中的變數數值和梯度數值，大家可以和前面推導的計算過程印證看看。

import torch
x1 = torch.tensor(2., requires_grad=True)
x2 = torch.tensor(5., requires_grad=True)

v1 = torch.log(x1)
v2 = x1 * x2
v3 = torch.sin(x2)
v4 = v1 + v2
y = v4 - v3

dv4 = torch.autograd.grad(y, [v4], retain_graph=True)
dv3 = torch.autograd.grad(y, [v3], retain_graph=True)
dv2 = torch.autograd.grad(v4, [v2], retain_graph=True)
dv1 = torch.autograd.grad(v4, [v1], retain_graph=True)
dv0 = torch.autograd.grad(v3, [x2], retain_graph=True)
dv21 = torch.autograd.grad(v2, [x1], retain_graph=True)
dv22 = torch.autograd.grad(v2, [x2], retain_graph=True)
dv11 = torch.autograd.grad(v1, [x1], retain_graph=True)

grads = torch.autograd.grad(y, [x1, x2])

print('v1 is :', v1)
print('v2 is :', v2)
print('v3 is :', v3)
print('y is :', y)

print('dv4 is :', dv4)
print('dv3 is :', dv3)
print('dv2 is :', dv2)
print('dv1 is :', dv1)
print('dv0 is :', dv0)
print('dv21 is :', dv21)
print('dv22 is :', dv22)
print('dv11 is :', dv11)

print('grad is : ', grads[0],grads[1])

輸出是

v1 is : tensor(0.6931, grad_fn=<LogBackward>)
v2 is : tensor(10., grad_fn=<MulBackward0>)
v3 is : tensor(-0.9589, grad_fn=<SinBackward>)
y is : tensor(11.6521, grad_fn=<SubBackward0>)
dv4 is : (tensor(1.),)
dv3 is : (tensor(-1.),)
dv2 is : (tensor(1.),)
dv1 is : (tensor(1.),)
dv0 is : (tensor(0.2837),)
dv21 is : (tensor(5.),)
dv22 is : (tensor(2.),)
dv11 is : (tensor(0.5000),)
grad is :  tensor(5.5000) tensor(1.7163)