torch.autograd 簡介

建立於：2017 年 3 月 24 日 | 最後更新：2025 年 1 月 10 日 | 最後驗證：2024 年 11 月 05 日

torch.autograd 是 PyTorch 的自動微分引擎，為神經網路訓練提供動力。在本節中，您將對 autograd 如何幫助神經網路訓練有一個概念性的了解。

背景

神經網路 (NN) 是在某些輸入資料上執行的一系列巢狀函數的集合。這些函數由參數（由權重和偏差組成）定義，這些參數在 PyTorch 中儲存在張量中。

NN 的訓練分兩個步驟進行

前向傳播：在前向傳播中，NN 會對正確的輸出做出最佳猜測。它會透過其每個函數執行輸入資料以進行此猜測。

反向傳播：在反向傳播中，NN 會根據其猜測中的錯誤來調整其參數。它透過從輸出向後遍歷，收集關於函數參數的誤差導數 (梯度)，並使用梯度下降來優化參數來完成此操作。有關反向傳播的更詳細的演練，請查看此 3Blue1Brown 的影片。

在 PyTorch 中的使用

讓我們看一下單個訓練步驟。在此範例中，我們從 torchvision 載入預訓練的 resnet18 模型。我們建立一個隨機資料張量，以表示具有 3 個通道、高度和寬度為 64 的單個影像，及其對應的 label，初始化為一些隨機值。預訓練模型中的 Label 具有形狀 (1,1000)。

注意

本教學課程僅適用於 CPU，不適用於 GPU 裝置 (即使張量已移至 CUDA)。

import torch
from torchvision.models import resnet18, ResNet18_Weights
model = resnet18(weights=ResNet18_Weights.DEFAULT)
data = torch.rand(1, 3, 64, 64)
labels = torch.rand(1, 1000)

Downloading: "https://download.pytorch.org/models/resnet18-f37072fd.pth" to /var/lib/ci-user/.cache/torch/hub/checkpoints/resnet18-f37072fd.pth

  0%|          | 0.00/44.7M [00:00<?, ?B/s]
 45%|####5     | 20.2M/44.7M [00:00<00:00, 212MB/s]
 92%|#########2| 41.1M/44.7M [00:00<00:00, 215MB/s]
100%|##########| 44.7M/44.7M [00:00<00:00, 215MB/s]

接下來，我們透過其每一層執行輸入資料的模型，以進行預測。這是前向傳遞。

prediction = model(data) # forward pass

我們使用模型的預測和對應的標籤來計算誤差 (loss)。下一步是透過網路反向傳播此誤差。當我們在誤差張量上呼叫 .backward() 時，反向傳播開始。然後，Autograd 會計算每個模型參數的梯度，並將其儲存在參數的 .grad 屬性中。

loss = (prediction - labels).sum()
loss.backward() # backward pass

接下來，我們載入一個優化器，在本例中，SGD 的學習率為 0.01，且動量為 0.9。我們在優化器中註冊模型的所有參數。

optim = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)

最後，我們呼叫 .step() 以啟動梯度下降。優化器會依儲存在 .grad 中的梯度調整每個參數。

optim.step() #gradient descent

此時，您擁有訓練神經網路所需的一切。以下各節詳細說明了 autograd 的工作原理 - 您可以隨時跳過它們。

---

Autograd 中的微分

讓我們看看 autograd 如何收集梯度。我們使用 requires_grad=True 建立兩個張量 a 和 b。這向 autograd 發出訊號，表示應該追蹤對它們的每個操作。

import torch

a = torch.tensor([2., 3.], requires_grad=True)
b = torch.tensor([6., 4.], requires_grad=True)

我們從 a 和 b 建立另一個張量 Q。

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

Q = 3*a**3 - b**2

假設 a 和 b 是 NN 的參數，而 Q 是誤差。在 NN 訓練中，我們想要關於參數的誤差梯度，即

\[\frac{\partial Q}{\partial a} = 9a^2 \]

\[\frac{\partial Q}{\partial b} = -2b \]

當我們在 Q 上呼叫 .backward() 時，autograd 會計算這些梯度，並將它們儲存在相應張量的 .grad 屬性中。

我們需要在 Q.backward() 中顯式傳遞一個 gradient 參數，因為它是一個向量。gradient 是一個形狀與 Q 相同的張量，它表示 Q 對自身的梯度，即

\[\frac{dQ}{dQ} = 1 \]

或者，我們也可以將 Q 聚合為一個純量，並隱式地呼叫 backward，例如 Q.sum().backward()。

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

梯度現在會儲存在 a.grad 和 b.grad 中

# check if collected gradients are correct
print(9*a**2 == a.grad)
print(-2*b == b.grad)

tensor([True, True])
tensor([True, True])

選讀 - 使用 autograd 的向量微積分

在數學上，如果你有一個向量值函數 \(\vec{y}=f(\vec{x})\)，則 \(\vec{y}\) 相對於 \(\vec{x}\) 的梯度是一個 Jacobian 矩陣 \(J\)

\[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)\]

一般來說，torch.autograd 是一個計算向量-Jacobian 乘積的引擎。也就是說，給定任何向量 \(\vec{v}\)，計算乘積 \(J^{T}\cdot \vec{v}\)

如果 \(\vec{v}\) 恰好是一個純量函數 \(l=g\left(\vec{y}\right)\) 的梯度

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

那麼根據鏈式法則，向量-Jacobian 乘積將是 \(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)\]

向量-Jacobian 乘積的這一特性就是我們在上面的例子中使用的；external_grad 代表 \(\vec{v}\)。

計算圖

從概念上講，autograd 會在由 Function 物件組成的有向無環圖 (DAG) 中，記錄資料（張量）和所有已執行的操作（以及產生的新張量）。在這個 DAG 中，葉節點是輸入張量，根節點是輸出張量。透過追蹤從根節點到葉節點的圖，你可以使用鏈式法則自動計算梯度。

在前向傳播中，autograd 同時做兩件事

• 執行請求的操作以計算產生的張量，以及

• 在 DAG 中維護操作的*梯度函數*。

當在 DAG 根節點上呼叫 .backward() 時，會觸發反向傳播。autograd 接著

• 從每個 .grad_fn 計算梯度，

• 將它們累積到相應張量的 .grad 屬性中，並且

• 使用鏈式法則，一路傳播到葉張量。

以下是我們範例中 DAG 的視覺表示。在圖中，箭頭方向是前向傳播的方向。節點代表前向傳播中每個操作的反向函數。藍色的葉節點代表我們的葉張量 a 和 b。

注意

DAG 在 PyTorch 中是動態的 需要注意的重要一點是，圖是從頭開始重新創建的；每次呼叫 .backward() 後，autograd 都會開始填入一個新的圖。這正是允許你在模型中使用控制流程語句的原因；你可以在每次迭代時根據需要更改形狀、大小和操作。

從 DAG 中排除

torch.autograd 會追蹤所有將 requires_grad 標誌設為 True 的張量上的操作。對於不需要梯度的張量，將此屬性設定為 False 會將其從梯度計算 DAG 中排除。

即使只有一個輸入張量具有 requires_grad=True，操作的輸出張量也會需要梯度。

x = torch.rand(5, 5)
y = torch.rand(5, 5)
z = torch.rand((5, 5), requires_grad=True)

a = x + y
print(f"Does `a` require gradients?: {a.requires_grad}")
b = x + z
print(f"Does `b` require gradients?: {b.requires_grad}")

Does `a` require gradients?: False
Does `b` require gradients?: True

在 NN 中，不計算梯度的參數通常稱為凍結參數。如果你事先知道不需要這些參數的梯度，那麼「凍結」模型的一部分會很有用（這可以透過減少 autograd 計算來提供一些效能優勢）。

在微調中，我們會凍結模型的大部分，並且通常只修改分類器層，以便對新的標籤進行預測。讓我們來看看一個小範例來演示這一點。和以前一樣，我們載入一個預訓練的 resnet18 模型，並凍結所有參數。

from torch import nn, optim

model = resnet18(weights=ResNet18_Weights.DEFAULT)

# Freeze all the parameters in the network
for param in model.parameters():
    param.requires_grad = False

假設我們想在一個具有 10 個標籤的新資料集上微調模型。在 resnet 中，分類器是最後一個線性層 model.fc。我們可以簡單地用一個新的線性層（預設為未凍結）來取代它，作為我們的分類器。

model.fc = nn.Linear(512, 10)

現在，除了 model.fc 的參數外，模型中的所有參數都被凍結。唯一計算梯度的參數是 model.fc 的權重和偏差。

# Optimize only the classifier
optimizer = optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)

請注意，儘管我們在最佳化器中註冊了所有參數，但唯一計算梯度（並因此在梯度下降中更新）的參數是分類器的權重和偏差。

同樣的排除功能也可以作為 torch.no_grad() 中的上下文管理器使用

---

延伸閱讀：

• 原地操作 & 多執行緒 Autograd

• 反向模式自動微分的範例實作

• 影片：PyTorch Autograd Explained - In-depth Tutorial