Pytorch autograd.grad與autograd.backward詳解

引言

平時在寫 Pytorch 訓練腳本時，都是下面這種無腦按步驟走：

outputs = model(inputs)		# 模型前向推理
optimizer.zero_grad()		# 清除累積梯度
loss.backward()				# 模型反向求導
optimizer.step()			# 模型參數更新

對用戶屏蔽底層自動微分的細節，使得用戶能夠根據簡單的幾個 API 將模型訓練起來。這對于初學者當然是極好的，也是 Pytorch 這幾年一躍成為最流行的深度學習框架的主要原因：易用性。

但是，我們有時需要深究自動微分的機制，比如元學習方法 MAML （參考 Pytorch 代碼）中，需要分別根據支持集和查詢集的梯度按照不同的策略更新模型參數。這時還是需要了解一些 Pytorch 框架的自動微分機制。幸運的是，Pytorch 關于這部分的框架設計也很清晰，在參考了幾個博客之后，筆者將自己的對 Pytorch 自動微分機制接口總結在這里。

注意只是自動微分機制的 Python 接口，而非底層實現。

背景知識

計算圖

當今主流深度學習框架的計算圖主要有兩種形式：靜態圖（TensoFlow 1.x、Caffe …）和動態圖（Pytorch …）。兩者的卻別簡單說來就是：靜態圖是在模型確定之后就先生成一張計算圖，然后每次對于不同的輸入樣本，都直接丟到計算圖中跑；而動態圖則是對于每次樣本輸入都重新構建一張計算圖。從它們的區別也可以感受到它們彼此最重要的優劣勢：靜態圖速度快但是不夠靈活，動態圖靈活但速度稍慢。

在今天，各個框架中動態圖與靜態圖的區分也沒有那么絕對了。比如 TensorFlow 2.0 已經采用動態圖，而 Pytorch 也可通過 scripting/tracing 轉換成 JIT torchscript 靜態圖。但這不是本文的重點，對深度學習框架計算圖感興趣可參考：機器學習系統:設計與實現計算圖。

我們要討論的是 Pytorch 的自動微分機制，Pytorch 中主要是動態圖，即計算圖的搭建和計算是同時的，對每次輸入都是重新建圖計算。在 Pytorch 的計算圖里有兩種元素：數據（tensor）和運算（operation）。

運算：包括了加減乘除、開方、冪指對、三角函數等可微分運算。
數據：在 Pytorch 中，數據的形式一般就是張量 torch.Tensor。

tensor

Pytorch 中 tensor 具有如下屬性：

requires_grad：是否需要求導
- 關于 requires_grad 屬性的默認值。自己定義的葉子節點默認為 False，而非葉子節點默認為 True，神經網絡中的權重默認為 True。判斷哪些節點是True/False 的一個原則就是從你需要求導的葉子節點到 loss 節點之間是一條可求導的通路，這條通路上的節點的 requires_grad 都應當是 True。
grad_fn：當前節點是經過什么運算（如加減乘除等）得到的
grad：導數值
data：tensor 的數據
is_leaf：是否為葉子節點
- 其他幾個概念都比較好理解，這里解釋一下什么是葉子節點。
- 在 Pytorch 中，如果一個張量的 requires_grad=True，則進一步可分為：葉子節點和非葉子節點。葉子節點是用戶創建的節點，不依賴其它節點，非葉子結點則是由葉子結點計算得到的中間張量。
```
a = torch.randn(2, 2).requires_grad_()
b = a * 2
print(a.is_leaf, b.is_leaf)
# 輸出：True False
```
- 對于 requires_grad=False 的 tensor 來說，我們約定俗成地把它們歸為葉子張量。但其實無論如何劃分都沒有影響，因為張量的 is_leaf 屬性只有在需要求導的時候才有意義。
- 由于葉子節點是用戶創建的，所以它的 grad_fn 為空，而非葉子節點都是經過運算得到的，所以 grad_fn 非空
- 葉子/非葉子表現出來的區別在于：反向傳播結束之后，非葉子節點的梯度會被釋放掉，只保留葉子節點的梯度，這樣就節省了內存。如果想要保留非葉子節點的梯度，可以使用 retain_grad() 方法。

關于 Pytorch tensor 的更多細節，可參考：淺談 PyTorch 中的 tensor 及使用。

一個例子

以下例子來自：PyTorch 的 Autograd。

了解過背景知識之后，現在我們來看一個具體的計算例子，先用最常見的梯度反傳方式 loss.backward() ，并畫出它的正向和反向計算圖。假如我們需要計算這么一個模型：

l1 = input x w1
l2 = l1 + w2
l3 = l1 x w3
l4 = l2 x l3
loss = mean(l4)

這個例子比較簡單，涉及的最復雜的操作是求平均，但是如果我們把其中的加法和乘法操作換成卷積，那么其實和神經網絡類似。我們可以簡單地畫一下它的計算圖，其中綠色節點表示葉子節點：

在這里插入圖片描述

圖1：正向計算圖

下面給出了對應的代碼，我們定義了 input，w1，w2，w3 這三個變量，其中 input 不需要求導結果。根據 Pytorch 默認的求導規則，對于 l1 來說，因為有一個輸入需要求導（也就是 w1 需要），所以它自己默認也需要求導，即 requires_grad=True（即前面提到的 ”是否在需要求導的通路上“ ，如果對這個規則不熟悉，歡迎參考淺談 PyTorch 中的 tensor 及使用或者直接查看官方 Tutorial 相關部分）。在整張計算圖中，只有 input 一個變量是 requires_grad=False 的。正向傳播過程的具體代碼如下：

input = torch.ones([2, 2], requires_grad=False)
w1 = torch.tensor(2.0, requires_grad=True)
w2 = torch.tensor(3.0, requires_grad=True)
w3 = torch.tensor(4.0, requires_grad=True)l1 = input * w1
l2 = l1 + w2
l3 = l1 * w3
l4 = l2 * l3
loss = l4.mean()print(w1.data, w1.grad, w1.grad_fn)
# tensor(2.) None Noneprint(l1.data, l1.grad, l1.grad_fn)
# tensor([[2., 2.],
#         [2., 2.]]) None <MulBackward0 object at 0x000001EBE79E6AC8>print(loss.data, loss.grad, loss.grad_fn)
# tensor(40.) None <MeanBackward0 object at 0x000001EBE79D8208>

正向傳播的結果基本符合我們的預期。我們可以看到，變量 l1 的 grad_fn 儲存著乘法操作符 <MulBackward0>，用于在反向傳播中指導導數的計算。而 w1 是用戶自己定義的，不是通過計算得來的，所以其 grad_fn 為空；同時因為還沒有進行反向傳播，grad 的值也為空。接下來，我們看一下如果要繼續進行反向傳播，計算圖應該是什么樣子：

在這里插入圖片描述

圖2：反向計算圖

反向圖也比較簡單，從 loss 這個變量開始，通過鏈式法則，依次計算出各部分的導數。說到這里，我們不妨先自己手動推導一下求導的結果，再與程序運行結果作對比。如果對這部分不感興趣的讀者，可以直接跳過。

再擺一下公式：

input = [1.0, 1.0, 1.0, 1.0]
w1 = [2.0, 2.0, 2.0, 2.0]
w2 = [3.0, 3.0, 3.0, 3.0]
w3 = [4.0, 4.0, 4.0, 4.0]l1 = input x w1 = [2.0, 2.0, 2.0, 2.0]
l2 = l1 + w2 = [5.0, 5.0, 5.0, 5.0]
l3 = l1 x w3 = [8.0, 8.0, 8.0, 8.0] 
l4 = l2 x l3 = [40.0, 40.0, 40.0, 40.0] 
loss = mean(l4) = 40.0

首先 $loss=14∑i=03l4iloss=\frac{1}{4}\sum_{i=0}^3l_4^i$ , 所以 $l oss$ 對 $l_4^i$ 的偏導分別為 $?loss?l4i=14\frac{\partial loss}{\partial l_4^i}=\frac{1}{4}$ ;

接著 $?l4?l3=l2=[5.0,5.0,5.0,5.0]\frac{\partial l_4}{\partial l_3}=l_2=[5.0,5.0,5.0,5.0]$ , 同時 $?l4?l2=l3=[8.0,8.0,8.0,8.0]\frac{\partial l_4}{\partial l_2}=l_3=[8.0,8.0,8.0,8.0]$ ;

現在看 $l_3$ 對它的兩個變量的偏導：

$?l3?l1=w3=[4.0,4.0,4.0,4.0]\frac{\partial l_3}{\partial l_1}=w3=[4.0,4.0,4.0,4.0]$ ， $?l3?w3=l1=[2.0,2.0,2.0,2.0]\frac{\partial l_3}{\partial w_3}=l1=[2.0,2.0,2.0,2.0]$

因此 $?loss?w3=?loss?l4?l4?l3?l3?w3=[2.5,2.5,2.5,2.5]\frac{\partial loss}{\partial w_3}=\frac{\partial loss}{\partial{l_4}}\frac{\partial{l_4}}{\partial{l_3}}\frac{\partial{l_3}}{\partial w_3}=[2.5,2.5,2.5,2.5]$ , 其和為 10 ;

同理，再看一下求 $w_2$ 導數的過程： $?loss?w2=?loss?l4?l4?l2?l2?w3=[2.0,2.0,2.0,2.0]\frac{\partial loss}{\partial w_2}=\frac{\partial loss}{\partial{l_4}}\frac{\partial{l_4}}{\partial{l_2}}\frac{\partial{l_2}}{\partial w_3}=[2.0,2.0,2.0,2.0]$ ，其和為 8。

其他的導數計算基本上都類似，因為過程太多，這里就不全寫出來了，如果有興趣的話大家不妨自己繼續算一下。

接下來我們繼續運行代碼，并檢查一下結果和自己算的是否一致：

loss.backward()print(w1.grad, w2.grad, w3.grad)
# tensor(28.) tensor(8.) tensor(10.)
print(l1.grad, l2.grad, l3.grad, l4.grad, loss.grad)
# None None None None None

首先我們需要注意一下的是，在之前寫程序的時候我們給定的 w 們都是一個常數，利用了廣播的機制實現和常數和矩陣的加法乘法，比如 w2 + l1，實際上我們的程序會自動把 w2 擴展成 [[3.0, 3.0], [3.0, 3.0]]，和 l1 的形狀一樣之后，再進行加法計算，計算的導數結果實際上為 [[2.0, 2.0], [2.0, 2.0]]，為了對應常數輸入，所以最后 w2 的梯度返回為矩陣之和 8 。

另外還有一個問題，注意到 l1，l2，l3，以及其他的部分的求導結果都為空。這驗證了我們之前提到的葉子結點的概念，對于非葉子幾點，不會保留其梯度值，如果一定要保留，需要設置 retain_graph=True。

torch.autograd：grad與backward

自動微分機制是深度學習框架的核心，對于 Pytorch 也不例外。 [Pytorch autograd官方文檔][https://pytorch.org/docs/stable/autograd.html#]指出，Pytorch 中有兩種方式可以實現反向傳播求導，分別是 torch.auograd.grad 和 torch.autograd.backward 。

在我們日常搭建訓練腳本的過程中，最常見的是 loss.backward() 。其實這是與 torch.autograd.backward(loss) 是等價的，即上述后一種方式。

兩種方式的區別是：前者是返回參數的梯度值列表，而后者是直接修改各個 tensor 的 grad 屬性。

接口定義

torch.autograd.backward

torch.autograd.backward(tensors, grad_tensors=None, retain_graph=None, create_graph=False, grad_variables=None)

tensor：用于計算梯度的tensor。前面提到過以下兩種方式是等價的：torch.autograd.backward(z) == z.backward()
grad_tensors：在計算矩陣的梯度時會用到。也是一個tensor，shape一般需要和前面的 tensor 保持一致。
retain_graph：通常在調用一次 grad/backward 后，Pytorch會自動把計算圖銷毀，所以要想對某個變量重復調用 backward，則需要將該參數設置為 True
create_graph：當設置為 True 的時候可以用來計算更高階的梯度
grad_variables：這個官方說法是 ‘grad_variables’ is deprecated. Use ‘grad_tensors’ instead.也就是說這個參數后面版本中應該會丟棄，直接使用 grad_tensors 就好了。

torch.autograd.grad

torch.autograd.grad(outputs, inputs, grad_outputs=None, retain_graph=None, create_graph=False, only_inputs=True, allow_unused=False)

outputs：結果節點，通常是損失值
inputs：需要求梯度的葉子節點，通常是模型參數
grad_outputs：類似于 backward 方法中的 grad_tensors
retain_graph：同上
create_graph：同上
only_inputs：默認為 True。如果為 True, 則只會返回指定 input 的梯度值。若為 False，則會計算所有葉子節點的梯度，并且將計算得到的梯度累加到各自的grad屬性上去。
allow_unused：默認為 False，即必須要指定input，如果沒有指定的話則報錯。

例子

還是通過一個例子來看：

import torch
import torch.nn as nnclass MyModel(nn.Module):def __init__(self):super().__init__()self.conv1 = nn.Conv2d(in_channels=2, out_channels=2, kernel_size=1, padding=0, bias=False)self.conv2 = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=1, padding=0, bias=False)def forward(self, z):return self.conv2(self.conv1(z))c = 2
h = 5
w = 5
lr = 0.01
inputs = torch.arange(0, c * h * w).float().view(1, c, h, w)
model = MyModel()
outputs = model(inputs)loss = outputs.sum()model.zero_grad()
grad = torch.autograd.grad(loss, model.parameters(), retain_graph=True)
# grad = torch.autograd.grad(loss, model.parameters())
# 注意這里需要 retain_grad = True，否則會報錯：
# RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.
loss.backward()for i, (name, param) in enumerate(model.named_parameters()):print("******************")print(name)print("grad using loss.backward: ", param.grad.data)print("grad using autograd.grad: ", grad[i])print("******************")# 更新參數# 相當于 optimizer.step()# theta_1 = theta_0 - lr * gradparam.data.sub_(lr * param.grad.data)# 或者：# param.data.sub_(lr * grad[i])

我們定義了一個簡單的兩層卷積模型，然后分別用 grad 和 backward 的方式來計算它們的梯度，并打印出來比較一下，發現是完全一致的。

如果想要根據梯度更新參數的話，也可以在拿到梯度之后，直接按照梯度下降的公式手動進行更新：
$θ1=θ0?α?θ0\theta_1=\theta_0-\alpha \nabla\theta_0$
這一步就相當于執行了 optimizer.step() ，它會使用封裝好的優化器進行更新。

求高階導

如何求高階導，比如求二階導，無非就是 grad_x 再對 x 求梯度：

x = torch.tensor(2.).requires_grad_()
y = torch.tensor(3.).requires_grad_()z = x * x * ygrad_x = torch.autograd.grad(outputs=z, inputs=x, retain_graph=True)
grad_xx = torch.autograd.grad(outputs=grad_x, inputs=x)print(grad_xx[0])
# 報錯：RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

報錯了，雖然 retain_graph=True 保留了計算圖和中間變量梯度，但沒有保存 grad_x 的運算方式，需要使用 create_graph=True 在保留原圖的基礎上再建立額外的求導計算圖，也就是會把 $?z?y=2xy\frac{\partial{z}}{\partial{y}}=2xy$ 這樣的運算存下來。

一階二階導我們可以分別用 autograd.grad 或者 backward 來做，即我們有四種排列組合，都是可以的：

# autograd.grad() + autograd.grad()
x = torch.tensor(2.).requires_grad_()
y = torch.tensor(3.).requires_grad_()z = x * x * ygrad_x = torch.autograd.grad(outputs=z, inputs=x, create_graph=True)
grad_xx = torch.autograd.grad(outputs=grad_x, inputs=x)print(grad_xx[0])
# 輸出：tensor(6.)

grad_xx 這里也可以直接用 backward，相當于直接從 $?z?y=2xy\frac{\partial{z}}{\partial{y}}=2xy$ 開始回傳

# autograd.grad() + backward()
x = torch.tensor(2.).requires_grad_()
y = torch.tensor(3.).requires_grad_()z = x * x * ygrad = torch.autograd.grad(outputs=z, inputs=x, create_graph=True)
grad[0].backward()print(x.grad)
# 輸出：tensor(6.)

也可以先用 backward 然后對 x.grad 這個一階導繼續求導

# backward() + autograd.grad()
x = torch.tensor(2.).requires_grad_()
y = torch.tensor(3.).requires_grad_()z = x * x * yz.backward(create_graph=True)
grad_xx = torch.autograd.grad(outputs=x.grad, inputs=x)print(grad_xx[0])
# 輸出：tensor(6.)

那是不是也可以直接用兩次 backward 呢？第二次直接 x.grad 從開始回傳，我們試一下

# backward() + backward()
x = torch.tensor(2.).requires_grad_()
y = torch.tensor(3.).requires_grad_()z = x * x * yz.backward(create_graph=True) # x.grad = 12
x.grad.backward()print(x.grad)
# 輸出：tensor(18., grad_fn=<CopyBackwards>)

發現了問題，結果不是 6，而是18，發現第一次回傳時輸出 x 梯度是12。這是因為 Pytorch 使用 backward 時默認會累加梯度，需要手動把前一次的梯度清零

x = torch.tensor(2.).requires_grad_()
y = torch.tensor(3.).requires_grad_()z = x * x * yz.backward(create_graph=True)
x.grad.data.zero_()
x.grad.backward()print(x.grad)
# 輸出：tensor(6., grad_fn=<CopyBackwards>)

對輸出矩陣自動微分

到此為止我們都是對標量進行自動微分，當我們試圖對向量或者矩陣進行梯度反傳時，會怎么樣呢？

import torch
x = torch.tensor([1., 2.]).requires_grad_()
y = x * xy.backward()
print(x.grad)
# 報錯：RuntimeError: grad can be implicitly created only for scalar outputs

報錯了，只有對標量輸出才能隱式地求梯度。即因為只能標量對標量，標量對向量求梯度， x 可以是標量或者向量，但 y 只能是標量；所以只需要先將 y 轉變為標量，對分別求導沒影響的就是求和。比如下面這樣：

import torch
x = torch.tensor([1., 2.]).requires_grad_()
y = x * xy = y.sum() # 求和，得到標量
y.backward()
print(x.grad)
# 輸出：tensor([2., 4.])

此時， $x=[x_1,y_1]$ ， $y=[x_1^2,x_2^2]$ ， $y'=y.sum()=x_1^2+x_2^2$ ，很顯然，求梯度有：
$?y′?x1=2x1=2?y′?x2=2x2=4\frac{\partial y'}{\partial x_1}=2x_1=2\ \ \ \ \ \ \ \ \ \ \frac{\partial y'}{\partial x_2}=2x_2=4$
與程序輸出相同。

為什么必須是標量呢？我們先寫出當輸出是一個向量 $y=[y_1,y_2]$ 時的雅克比矩陣：
$J=[?y?x1,?y?x2]=[?y1?x1?y1?x2?y2?x1?y2?x2]{J}=[\frac{\partial y}{\partial x_1},\frac{\partial y}{\partial x_2}]=\begin{bmatrix}{\frac{\partial y_1}{\partial x_1}}&{\frac{\partial y_1}{\partial x_2}}\\{\frac{\partial y_2}{\partial x_1}}&{\frac{\partial y_2}{\partial x_2}}\end{bmatrix}$
而我們想要的是 $[?y1?x1,?y2?x2][\frac{\partial y_1}{\partial x_1},\frac{\partial y_2}{\partial x_2}]$ ，從矩陣計算的角度來看，是不是只要對雅克比矩陣左乘個 $[1, 1]$ 就可以得到我們想要的了：
$[?y1?x1,?y2?x2]=[1,1]?J[\frac{\partial y_1}{\partial x_1},\frac{\partial y_2}{\partial x_2}]=[1,1]\cdot J$
這就是不使用 y.sum() 的另一種方式，通過 backward 接口的 grad_tensors 參數（上面介紹過）：

x = torch.tensor([1., 2.]).requires_grad_()
y = x * xy.backward(torch.ones_like(y))
print(x.grad)
# 輸出：tensor([2., 4.])

如果要使用 torch.autograd.grad ，對應的接口形參是 grad_outputs ：

x = torch.tensor([1., 2.]).requires_grad_()
y = x * xgrad_x = torch.autograd.grad(outputs=y, inputs=x, grad_outputs=torch.ones_like(y))
# 或者
# grad_x = torch.autograd.grad(outputs=y.sum(), inputs=x)
print(grad_x[0])
# 輸出：tensor([2., 4.])

實際上，grad_tensors 的作用其實可以簡單地理解成在求梯度時的權重，因為可能不同值的梯度對結果影響程度不同，所以 Pytorch 弄了個這種接口，而沒有固定為全是1。引用自知乎上的一個評論：如果從最后一個節點(總loss)來backward，這種實現(torch.sum(y*w))的意義就具體化為 multiple loss term with difference weights 這種需求了吧。

關于對輸出矩陣求微分，更詳細的可參考：PyTorch 的 backward 為什么有一個 grad_variables 參數？

幾個細節

zero_grad

在寫訓練腳本時，我們通常在每次 backward 反傳之前，都要進行一步 optimizer.zero_gard() ，這一步是做什么的呢？實際上就如同名字顯示那樣，本步的目的就是將目前葉子結點中上一步的梯度 grad 清零，然后再進行反傳，計算本 batch 的梯度。

那能不能不每次都清零梯度呢？實際上是可以的，這可以作為一種變相增大 batch size 的 trick。如果我們的機器每個 batch 最多只能 64 個樣本，那我們設置每步都計算梯度并累計到葉子結點的 grad 屬性中，但是每隔一步才進行一次參數更新和梯度清零，這就相當于 batch_size 成了 128。但這也會出現一些問題，比如 BN 怎么辦，這在知乎上也有一些問題有討論過，感興趣可以查一下。

model.zero_grad()還是optimizer.zero_grad()?

看代碼時，有時候會看到 model.zero_grad() ，有時又會看到 optimizer.zero_grad() ，到底有什么區別呢？

我們知道模型就是一堆參數按照特定的運算結構組織起來，我們在構建 optimizer 時會把優化器要優化的參數傳遞給它，比如：

optimizer = Adam(model.parameters(), lr=lr)

常規情況下傳入優化器的只有 model.parameters()，但是并不總是如此。有時候，整個模型要優化的不只有模型本身的參數，還可能有一些自定義的 parameters，比如：

pref_vec = torch.nn.Parameter(torch.randn(1, 512))
optimizer = Adam([{'params': model.parameters()}, {'params': pref_vec}], lr=lr)

在這種情況下 model.parameters() 與 pref_vec 是一起更新的，都有 optimizer 這個優化器來更新。

指出這一點之后，大家應該就明白 model.zero_grad() 和 optimizer.zero_grad() 的區別了。它們指向的待更新參數（葉子結點）不一定是一樣的。一般情況下（優化器待更新參數就是模型參數）二者是等價的，但是如果待更新的參數除了模型的參數之外還有一些自定義的參數，就必須用 optimizer.zero_grad() 了。

detach

detach 會切斷當前張量與計算圖之間的聯系，不會再往后計算梯度。

假設有模型 A 和模型 B，我們需要將 A 的輸出作為 B 的輸入，但訓練時我們只訓練模型B，那么可以這樣做：

input_B = output_A.detach()

inplace

inplace 操作顧名思義，就是直接在原地改變張量的值，而不是計算后得到一個新的張量并返回。

注意：葉子節點不可執行 in-place 操作，因為反向傳播時會訪問原來的對象地址。

關于 inplace 操作也有很多坑，經常見到的一個報錯是：

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: ...

關于 inplace 操作的問題在 PyTorch 的 Autograd 中有詳細的討論。

Ref

Pytorch autograd官方文檔
一文解釋PyTorch求導相關 (backward, autograd.grad)
MAML-Pytorch
機器學習系統:設計與實現計算圖
淺談 PyTorch 中的 tensor 及使用
PyTorch 的 Autograd
一文解釋PyTorch求導相關 (backward, autograd.grad)
PyTorch 的 backward 為什么有一個 grad_variables 參數？
Pytorch autograd,backward詳解