[源碼解析]深度學(xué)習(xí)利器之自動微分(3) --- 示例解讀

本文從 PyTorch 兩篇官方文檔開始為大家解讀兩個示例。本文不會逐句翻譯，而是選取重點(diǎn)并且試圖加入自己的理解。

[源碼解析]深度學(xué)習(xí)利器之自動微分(3) --- 示例解讀

[源碼解析]深度學(xué)習(xí)利器之自動微分(3) --- 示例解讀

0x00 摘要

本文從 PyTorch 兩篇官方文檔開始為大家解讀兩個示例。本文不會逐句翻譯，而是選取重點(diǎn)并且試圖加入自己的理解。

我們在前兩篇文章學(xué)習(xí)了自動微分的基本概念，從本文開始，我們繼續(xù)分析 PyTorch 如何實現(xiàn)自動微分。因為涉及內(nèi)容太多太復(fù)雜，所以計劃使用 2～3篇來介紹前向傳播如何實現(xiàn)，用 3 ～ 4 篇來介紹后向傳播如何實現(xiàn)。

系列前兩篇連接如下：

深度學(xué)習(xí)利器之自動微分(1)

深度學(xué)習(xí)利器之自動微分(2)

0x01 概述

在訓(xùn)練神經(jīng)網(wǎng)絡(luò)時，最常用的算法是 反向傳播。在該算法中根據(jù)損失函數(shù)相對于給定參數(shù)的梯度來對參數(shù)（模型權(quán)重）進(jìn)行調(diào)整。為了計算這些梯度，PyTorch 實現(xiàn)了一個名為 torch.autograd的內(nèi)置反向自動微分引擎。它支持任何計算圖的梯度自動計算。

1.1 編碼歷史

從概念上講，autograd 記錄了一個計算圖。在創(chuàng)建張量時，如果設(shè)置 requires_grad 為Ture，那么 Pytorch 就知道需要對該張量進(jìn)行自動求導(dǎo)。于是PyTorch會記錄對該張量的每一步操作歷史，從而生成一個概念上的有向無環(huán)圖，該無環(huán)圖的葉子節(jié)點(diǎn)是模型的輸入張量，其根為模型的輸出張量。用戶不需要對圖的所有執(zhí)行路徑進(jìn)行編碼，因為用戶運(yùn)行的就是用戶后來想微分的。通過從根到葉跟蹤此圖形，用戶可以使用鏈?zhǔn)角髮?dǎo)規(guī)則來自動計算梯度。

在內(nèi)部實現(xiàn)上看，autograd 將此圖表示為一個“Function” 或者說是"Node" 對象（真正的表達(dá)式）的圖，該圖可以使用apply方法來進(jìn)行求值。

1.2 如何應(yīng)用

在前向傳播計算時，autograd做如下操作：

運(yùn)行請求的操作以計算結(jié)果張量。
建立一個計算梯度的DAG圖，在DAG圖中維護(hù)所有已執(zhí)行操作（包括操作的梯度函數(shù)以及由此產(chǎn)生的新張量）的記錄。每個tensor梯度計算的具體方法存放于tensor節(jié)點(diǎn)的grad_fn屬性中。

當(dāng)向前傳播完成之后，我們通過在在 DAG 根上調(diào)用.backward() 來執(zhí)行后向傳播，autograd會做如下操作：

利用.grad_fn計算每個張量的梯度，并且依據(jù)此構(gòu)建出包含梯度計算方法的反向傳播計算圖。
將梯度累積在各自的張量.grad屬性中，并且使用鏈?zhǔn)椒▌t，一直傳播到葉張量。
每次迭代都會重新創(chuàng)建計算圖，這使得我們可以使用Python代碼在每次迭代中更改計算圖的形狀和大小。

需要注意是，PyTorch 中的DAG 是動態(tài)的，每次 .backward()調(diào)用后，autograd 開始填充新計算圖，該圖是從頭開始重新創(chuàng)建。這使得我們可以使用Python代碼在每次迭代中更改計算圖的形狀和大小。

0x02 示例

下面我們通過兩個例子來進(jìn)行解讀，之所以使用兩個例子，因為均來自于PyTorch 官方文檔。

2.2 實例解讀 1

我們首先使用 https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html 來進(jìn)行演示和解讀。

2.2.1 代碼

示例代碼如下：

import torch

a = torch.tensor(2., requires_grad=True)
b = torch.tensor(6., requires_grad=True)
O = 3*a**3
P = b**2
Q = O - P
external_grad = torch.tensor(1.)
Q.backward(gradient=external_grad)
print(a.grad)
print(b.grad)

print("=========== grad")

a = torch.tensor(2., requires_grad=True)
b = torch.tensor(6., requires_grad=True)
Q = 3*a**3 - b**2
grads = torch.autograd.grad(Q, [a, b])
print(grads[0])
print(grads[1])

print(Q.grad_fn.next_functions)
print(O.grad_fn.next_functions)
print(P.grad_fn.next_functions)
print(a.grad_fn)
print(b.grad_fn)

輸出為：

tensor(36.)
tensor(-12.)
=========== grad
tensor(36.)
tensor(-12.)

((<MulBackward0 object at 0x000001374DE6C308>, 0), (<PowBackward0 object at 0x000001374DE6C288>, 0))
((<PowBackward0 object at 0x000001374DE6C288>, 0), (None, 0))
((<AccumulateGrad object at 0x000001374DE6C6C8>, 0),)
None
None

這里的Q運(yùn)算方式如下：

\[Q = 3a^3 - b^2 \]

因此Q對a, b 的求導(dǎo)如下：

\[\frac{?Q}{?a} = 9a^2 \\\frac{?Q}{?b} = -2b \]

2.2.2 分析

動態(tài)圖是在前向傳播的時候建立。前向傳播時候，Q是最終的輸出，但是在反向傳播的時候，Q 卻是計算的最初輸入，就是反向傳播圖的Root。

示例中，對應(yīng)的張量是：

a 是 2，b 是 6， Q 是 tensor(-12., grad_fn=<SubBackward0>)。

對應(yīng)的積分是：

Q對于 a 的積分是：$\frac{?Q}{?a} = 9a^2$ = 36。
Q對于b的積分是 $\frac{?Q}{?b} = -2b$ = -12。

當(dāng)我們調(diào)用.backward()時，backward()只是通過將其參數(shù)傳遞給已經(jīng)生成的反向圖來計算梯度。autograd 計算這些梯度并將它們存儲在各自的張量.grad屬性中。

我們需要顯式地給Q.backward()傳入一個gradient參數(shù)，因為它是一個向量。 gradient是與形狀相同的張量Q，它表示 Q 本身的梯度，即

\[\frac{?Q}{?Q} = 1 \]

等效地，我們也可以將 Q 聚合為一個標(biāo)量并隱式地向后調(diào)用，例如Q.sum().backward()。

external_grad = torch.tensor([1., 1.])
Q.backward(gradient=external_grad)

下面是我們示例中 DAG 的可視化表示。在圖中，箭頭指向前向傳遞的方向。節(jié)點(diǎn)代表前向傳遞中每個操作的后向函數(shù)。藍(lán)色的葉子節(jié)點(diǎn)代表我們的葉子張量a和b。

2.3 實例解讀 2

這次以https://pytorch.org/tutorials/beginner/basics/autogradqs_tutorial.html為例子說明。

2.3.1 示例代碼

考慮最簡單的一層神經(jīng)網(wǎng)絡(luò)，具有輸入x、參數(shù)w和b，以及一些損失函數(shù)。它可以通過以下方式在 PyTorch 中定義：

import torch

x = torch.ones(5)  # input tensor
y = torch.zeros(3)  # expected output
w = torch.randn(5, 3, requires_grad=True)
b = torch.randn(3, requires_grad=True)
z = torch.matmul(x, w)+b
loss = torch.nn.functional.binary_cross_entropy_with_logits(z, y)

2.3.2 張量、函數(shù)和計算圖

上述代碼定義了以下計算圖：

圖片來源是：https://pytorch.org/tutorials/_images/comp-graph.png

在這個網(wǎng)絡(luò)中，w和b是我們需要優(yōu)化的參數(shù)。因此，我們需要計算關(guān)于這些變量的損失函數(shù)的梯度。為了做到這一點(diǎn)，我們設(shè)置了這些張量的requires_grad屬性。

注意，您可以在創(chuàng)建張量時設(shè)置requires_grad的值，也可以稍后使用x.requires_grad_(True)方法設(shè)置。

我們應(yīng)用于張量來構(gòu)建計算圖的函數(shù)實際上是一個Function類的對象。該對象知道如何在前向計算函數(shù)，以及如何在反向傳播步驟中計算其導(dǎo)數(shù)。對反向傳播函數(shù)的引用存儲在grad_fn張量的屬性中。

print('Gradient function for z =', z.grad_fn)
print('Gradient function for loss =', loss.grad_fn)

輸出如下：

Gradient function for z = <AddBackward0 object at 0x7f4dbd4d3080>
Gradient function for loss = <BinaryCrossEntropyWithLogitsBackward object at 0x7f4dbd4d3080>

2.3.3 計算梯度

為了優(yōu)化神經(jīng)網(wǎng)絡(luò)中參數(shù)的權(quán)重，我們需要計算損失函數(shù)關(guān)于參數(shù)的導(dǎo)數(shù)，即我們需要在限定一些 x和y時候得到 $ \frac{\partial loss}{\partial w}$ 和 $\frac{\partial loss}{\partial b} $ 。為了計算這些導(dǎo)數(shù)，我們調(diào)用 loss.backward()，然后從w.grad和 b.grad 之中獲得數(shù)值：

loss.backward()
print(w.grad)
print(b.grad)

得出：

tensor([[0.1881, 0.1876, 0.0229],
        [0.1881, 0.1876, 0.0229],
        [0.1881, 0.1876, 0.0229],
        [0.1881, 0.1876, 0.0229],
        [0.1881, 0.1876, 0.0229]])
tensor([0.1881, 0.1876, 0.0229])

注意

我們只能獲取在計算圖葉子節(jié)點(diǎn)的requires_grad屬性設(shè)置為True時候得到該節(jié)點(diǎn)的grad屬性。我們沒法得到們圖中的所有其他節(jié)點(diǎn)的梯度。
出于性能原因，我們只能在給定的計算圖之上使用backward執(zhí)行一次梯度計算。如果我們需要在同一個圖上多次調(diào)用backward，則需要在backward調(diào)用時候設(shè)置 retain_graph=True。

2.3.4 禁用梯度跟蹤

默認(rèn)情況下，所有設(shè)置requires_grad=True 的張量都會跟蹤其計算歷史并支持梯度計算。但是，有些情況下我們不需要這樣做，例如，當(dāng)我們已經(jīng)訓(xùn)練了模型并且只想將其應(yīng)用于某些輸入數(shù)據(jù)時，即我們只想通過網(wǎng)絡(luò)進(jìn)行前向計算，這時候我們可以通過用torch.no_grad()塊包圍我們的計算代碼以停止跟蹤計算：

z = torch.matmul(x, w)+b
print(z.requires_grad)

with torch.no_grad():
    z = torch.matmul(x, w)+b
print(z.requires_grad)

輸出：

True
False

實現(xiàn)相同結(jié)果的另一種方法是在張量上使用detach()方法：

z = torch.matmul(x, w)+b
z_det = z.detach()
print(z_det.requires_grad)

輸出：

False

您可能想要禁用梯度跟蹤的原因有：

將神經(jīng)網(wǎng)絡(luò)中的某些參數(shù)標(biāo)記為凍結(jié)參數(shù)。這是微調(diào)預(yù)訓(xùn)練網(wǎng)絡(luò)的一個非常常見的場景。
在僅進(jìn)行前向傳遞時加快計算速度，因為對不跟蹤梯度的張量進(jìn)行計算會更有效。

0x03 邏輯關(guān)系

如果從計算圖角度來看前向計算的過程，就是在構(gòu)建圖和執(zhí)行圖。"構(gòu)建圖"描述的是節(jié)點(diǎn)運(yùn)算之間的關(guān)系。"執(zhí)行圖"則是在會話中執(zhí)行這個運(yùn)算關(guān)系，就是張量在計算圖之中進(jìn)行前向傳播的過程。

前向計算依賴一些基礎(chǔ)類，在具體分析前向傳播之前，我們先要看看這些基礎(chǔ)類之間的邏輯關(guān)系。從DAG角度來分析 PyTorch 這個系統(tǒng)，其具體邏輯如下。

圖表示計算任務(wù)。PyTorch把計算都當(dāng)作是一種有向無環(huán)圖，或者說是計算圖，但這是一種虛擬的圖，代碼中沒有真實的數(shù)據(jù)結(jié)構(gòu)。
計算圖由節(jié)點(diǎn)（Node）和邊（Edge）組成。
節(jié)點(diǎn)（Node）代表了運(yùn)算操作。
- 一個節(jié)點(diǎn)通過邊來獲得 0 個或多個 Tensor，節(jié)點(diǎn)執(zhí)行計算之后會產(chǎn)生 0 個或多個 Tensor。
- 節(jié)點(diǎn)的成員變量 next_functions 是一個 tuple 列表，此列表就代表本節(jié)點(diǎn)要輸出到哪些其他 Function。列表個數(shù)就是這個 grad_fn 的 Edge 數(shù)目，列表之中每一個 tuple 對應(yīng)一條 Edge 信息，內(nèi)容就是 (Edge.function, Edge.input_nr)。
邊（Edge）就是運(yùn)算操作之間的流向關(guān)系。
- Edge.function ：表示此 Edge 需要輸出到哪一個其他 Function。
- Edge.input_nr ：指定本 Edge 是 Function 的第幾個輸入。
使用張量（ Tensor） 表示數(shù)據(jù)，就是在節(jié)點(diǎn)間流動的數(shù)據(jù)，如果沒有數(shù)據(jù)，計算圖就沒有任何意義。

具體可以參見下圖：

+---------------------+              +----------------------+
| SubBackward0        |              | PowBackward0         |
|                     |      Edge    |                      |  Edge
|   next_functions  +-----+--------> |     next_functions +----------> ...
|                     |   |          |                      |
+---------------------+   |          +----------------------+
                          |
                          |
                          |          +----------------------+
                          |  Edge    | MulBackward0         |
                          +--------> |                      |  Edge
                                     |     next_functions +----------> ...
                                     |                      |
                                     +----------------------+

至此，示例解析結(jié)束，我們下一篇介紹PyTorch 微分引擎相關(guān)的一些基礎(chǔ)類。