PyTorch 訓練迴圈：zero_grad、backward、step 為什麼是這個順序

💡 先搞懂問題

虛構的「北辰醫院」資訊室工程師阿哲，照著網路範例寫了一段 PyTorch 程式，用病房感測器資料預測病人是否需要加強巡視。他覺得每一批都呼叫 optimizer.zero_grad() 有點多餘，就把那一行刪掉，結果 loss 先降後升、來回震盪，怎麼調學習率都不穩。同組的另一位同事程式沒刪任何東西，但同一份驗證資料每跑一次準確率就不一樣，原因是驗證前忘了 model.eval()。兩支程式都沒有任何錯誤訊息。

scikit-learn 的 fit 把訓練過程藏起來，PyTorch 則要你自己寫訓練迴圈（training loop），所以每一行的先後都有意義。一個 batch（一小批資料）的更新固定是五個動作：前向傳播（forward，用目前的參數算出預測）、算損失（loss，預測和答案差多少）、清掉舊梯度（zero_grad）、反向傳播（backward，由自動微分 autograd 算出 loss 對每個參數的梯度 gradient，存進參數的 .grad），最後由最佳化器（optimizer）step 依梯度更新參數。新手最常卡在一個設計：.grad 預設是「加上去」而不是「蓋過去」，所以不清的話，每一批的梯度都會疊在前面所有批次上。

生活比喻：射箭社的練習

晴空射箭社的學員小芸每天練習五十輪。每一輪的流程是：先換上一張新靶紙，射出一箭，量箭落點離靶心多遠，教練看著彈孔分析「手腕往左偏了一點、站姿往前傾太多」，小芸再依照分析稍微調整姿勢。調整幅度不能太大，教練通常只要她改一點點，否則容易矯枉過正。

有一天助教忘了換靶紙。第二輪教練看到靶上有兩個彈孔，以為兩個都是這一輪射的，就把兩次的偏差加起來分析，要小芸大幅修正；第三輪靶上有三個彈孔，修正得更誇張，小芸的箭開始在靶心兩側來回亂跳。另外，平常練習時教練會隨機要她拿掉某個輔助動作，例如這一輪不准看瞄準器，逼她不要依賴單一線索；到了正式比賽，所有動作都恢復，而且比賽中不做分析、只記成績。

比喻有三個地方和實際不同。第一，射箭一次只調一個人的姿勢，模型一次 step 會同時調整所有參數，動輒上百萬個。第二，梯度不是「偏離靶心的位置」，而是「參數往正方向動一點，loss 會增加多少」，指向 loss 變大的方向，所以 step 是往反方向走：w ← w − 學習率 × w.grad（這是 SGD 的規則，Adam 等最佳化器會再做調整）。第三，梯度累加在數學上就是單純相加，PyTorch 這樣設計是刻意的，可以拿來模擬更大的 batch；問題只在於「不知不覺地」累加。

🎮 互動實驗室一：逐步執行器

模型只有一個參數 w，預測值是 w·x。示意資料 3 筆：x = 1、2、3，y = 2、4、6，所以理想的 w 是 2。損失用 MSE（均方誤差），最佳化器是學習率 0.1 的 SGD，3 筆資料整份當成一個 batch，所以每一輪（epoch）只更新一次。按「下一步」執行反白的那一行，右邊會顯示 w、pred、loss、w.grad 目前的值；上方可以改變 zero_grad 的位置，或打開「忘記寫 zero_grad」，看梯度怎麼累加、w 怎麼走歪。畫面上的數字都和 PyTorch 2.14 實跑結果一致（四捨五入到小數第 4 位）。

x = torch.tensor([1., 2., 3.])
y = torch.tensor([2., 4., 6.])
w = torch.tensor(0.0, requires_grad=True)
opt = torch.optim.SGD([w], lr=0.1)
for epoch in range(8):
    opt.zero_grad()                  # 清梯度
    pred = w * x                     # 前向
    loss = ((pred - y) ** 2).mean()  # 算 loss
    loss.backward()                  # 反向
    opt.step()                       # 更新

x = torch.tensor([1., 2., 3.])
y = torch.tensor([2., 4., 6.])
w = torch.tensor(0.0, requires_grad=True)
opt = torch.optim.SGD([w], lr=0.1)
for epoch in range(8):
    pred = w * x                     # 前向
    loss = ((pred - y) ** 2).mean()  # 算 loss
    opt.zero_grad()                  # 清梯度
    loss.backward()                  # 反向
    opt.step()                       # 更新

x = torch.tensor([1., 2., 3.])
y = torch.tensor([2., 4., 6.])
w = torch.tensor(0.0, requires_grad=True)
opt = torch.optim.SGD([w], lr=0.1)
for epoch in range(8):
    pred = w * x                     # 前向
    loss = ((pred - y) ** 2).mean()  # 算 loss
    loss.backward()                  # 反向
    opt.step()                       # 更新
    opt.zero_grad()                  # 清梯度（給下一輪）

x = torch.tensor([1., 2., 3.])
y = torch.tensor([2., 4., 6.])
w = torch.tensor(0.0, requires_grad=True)
opt = torch.optim.SGD([w], lr=0.1)
for epoch in range(8):
    # opt.zero_grad()                ← 這行忘了寫
    pred = w * x                     # 前向
    loss = ((pred - y) ** 2).mean()  # 算 loss
    loss.backward()                  # 反向
    opt.step()                       # 更新

🎮 互動實驗室二：步驟排序

三回合由小到大：第一回合排一個 batch 的五步，第二回合排一個 epoch 的訓練加驗證，第三回合排 Keras 的寫法做對照。把左邊的步驟卡排進右邊的編號格：電腦上可以直接拖曳；手機或想用點的，先點一張卡（會被框起來）再點格子。點選格子裡的卡片後再點另一格可以互換，選取後點左邊卡片區可以退回。排滿後按「檢查順序」，排錯的格子會晃動，並說明那樣寫實際會發生什麼。正確答案不只一種。

for xb, yb in train_loader:
    opt.zero_grad()            # 清梯度（也可放在 loss 之後，或 step 之後）
    pred = model(xb)           # 前向
    loss = loss_fn(pred, yb)   # 算 loss
    loss.backward()            # 反向：梯度加進 .grad
    opt.step()                 # 依 .grad 更新參數

model.train()                          # 1 切回訓練模式
for xb, yb in train_loader:            # 2 每一批做五步
    opt.zero_grad(); loss = loss_fn(model(xb), yb)
    loss.backward(); opt.step()
model.eval()                           # 3 Dropout、BatchNorm 改用評估行為
with torch.no_grad():                  # 4 不建計算圖
    val_loss = 0.0
    for xb, yb in val_loader:          # 5 驗證迴圈
        val_loss += loss_fn(model(xb), yb).item()
print(val_loss / len(val_loader))      # 6 平均驗證 loss

model = keras.Sequential([keras.Input(shape=(8,)),
                          layers.Dense(32, activation="relu"),
                          layers.Dense(2, activation="softmax")])
model.compile(optimizer="adam",                         # 對應 torch.optim
              loss="sparse_categorical_crossentropy",   # 對應 loss_fn
              metrics=["accuracy"])
model.fit(X_tr, y_tr, epochs=5, batch_size=32)          # 內含五步的雙層迴圈
model.evaluate(X_te, y_te)                              # 對應 eval＋no_grad
model.predict(X_new)

🎮 互動實驗室三：這段訓練程式哪裡有問題

每張卡是一段原創的 PyTorch 程式片段（model、loss_fn、opt、loader 都已建立好）。請判斷它的問題屬於哪一類，或者根本沒有問題。有幾張會直接報錯，更多張是「能跑但結果不對」。答完會說明實際會發生什麼、怎麼改，旁邊的分類表會亮起正解。

📘 原理補完

1. 五步各自讀什麼、改什麼

把每一步「需要什麼」與「改了什麼」寫下來，順序就不必死背：loss 需要前向算出的預測；backward 需要 loss 和它背後的計算圖（computational graph）；step 需要 backward 寫進 .grad 的梯度。唯一位置有彈性的是 zero_grad，它只負責在這一批的 backward 之前把舊梯度清掉。

一個 epoch 是把整份訓練資料看過一遍；每一批（batch）做一次上面的五步，稱為一次 iteration。所以一個 epoch 裡 step 的次數等於批次數，也就是 len(train_loader)，在 drop_last=False（預設）時是「樣本數 ÷ batch_size 無條件進位」。

2. 梯度為什麼會累加

autograd 在 backward 時做的是 param.grad += 新梯度，而不是直接覆蓋。這樣設計有實際用途：記憶體放不下大 batch 時，可以連續跑幾個小 batch 各自 backward，讓梯度自然加總，再 step 一次（梯度累積，gradient accumulation）；一個模型有好幾個 loss 時，也可以分別 backward 再一起更新。代價是一般訓練必須每批清一次。PyTorch 2.0 起 zero_grad() 預設 set_to_none=True，把 .grad 設成 None 而不是填 0，省一次記憶體寫入；對 SGD、Adam 這些內建最佳化器來說，.grad 是 None 的參數在 step 時會直接被略過。

下面這段程式就是實驗室一的完整版，可以直接執行，註解裡的輸出是 PyTorch 2.14.1 的實跑結果：

import torch

def train(forget_zero_grad):
    x = torch.tensor([1., 2., 3.])
    y = torch.tensor([2., 4., 6.])               # 示意資料：y 剛好是 x 的 2 倍
    w = torch.tensor(0.0, requires_grad=True)    # 唯一的參數，從 0 開始
    opt = torch.optim.SGD([w], lr=0.1)
    history = []
    for epoch in range(6):
        if not forget_zero_grad:
            opt.zero_grad()                      # 1. 清掉上一輪的梯度（設成 None）
        pred = w * x                             # 2. 前向
        loss = ((pred - y) ** 2).mean()          # 3. 算 loss（MSE）
        loss.backward()                          # 4. 反向：梯度「加」進 w.grad
        opt.step()                               # 5. 更新：w = w - 0.1 * w.grad
        history.append(round(w.item(), 4))       # .item() 取出 Python 數字
    return history

print(train(forget_zero_grad=False))  # [1.8667, 1.9911, 1.9994, 2.0, 2.0, 2.0]
print(train(forget_zero_grad=True))   # [1.8667, 3.8578, 4.115, 2.3982, 0.3098, -0.2011]

3. backward 在算什麼

前向傳播時，只要參與運算的張量有 requires_grad=True（nn.Module 的參數預設就是），PyTorch 就會記下每一步運算，形成計算圖。對純量 loss 呼叫 backward()，autograd 沿著這張圖反向套用連鎖律，算出 loss 對每個葉節點參數的偏微分。以示意模型在 w = 0 時為例：loss = mean((w·x − y)²)，對 w 的偏微分是 mean(2·(w·x − y)·x) = (2·(−2)·1 + 2·(−4)·2 + 2·(−6)·3) ÷ 3 = −18.6667，和 PyTorch 算出的 w.grad 相同。backward 完成後，計算圖預設會被釋放，所以同一個 loss 不能 backward 兩次（除非指定 retain_graph=True）。

4. model.train()、model.eval() 與 torch.no_grad()

這三個常被混在一起，其實管的是兩件不同的事。model.train() 與 model.eval() 切換的是「模式」：它們把模型與所有子層的 training 旗標設成 True 或 False，影響 Dropout（訓練時隨機把部分輸出設成 0，並把其餘放大 1/(1−p) 倍；評估時什麼都不做）與 BatchNorm（訓練時用這一批的平均與變異數並更新累計值，評估時改用累計值）。它們不會訓練、也不會關掉梯度。torch.no_grad() 管的是 autograd：區塊裡的運算不建計算圖，輸出的 requires_grad 是 False，省記憶體也比較快，但不影響 Dropout。驗證與推論時兩個都要：model.eval() 讓輸出正確且穩定，torch.no_grad()（或更嚴格的 torch.inference_mode()）讓它省資源。

5. loss.item() 與越吃越多的記憶體

每一批的 loss 是一個帶著 grad_fn 的張量，它和產生它的那張計算圖連在一起。如果為了算平均而寫 total += loss，total 本身也變成一個帶 grad_fn 的張量（實跑 total.requires_grad 是 True），它會一路參照每一批的 loss 與相關的計算圖，整個 epoch 的這些物件都無法被回收，記憶體隨批次增加，這是 PyTorch 官方 FAQ 特別提醒的寫法。改成 total += loss.item()，取出的是單純的 Python 浮點數，和計算圖完全脫鉤。.item() 只能用在只有一個元素的張量上；要把多個元素的張量轉成 NumPy，則要先 .detach()（在 GPU 上還要 .cpu()）。

6. 和 Keras 的 compile／fit 對照

Keras 把整個雙層迴圈包進 fit：compile 負責指定最佳化器與損失函數，相當於 PyTorch 建立 opt 與 loss_fn；fit(epochs, batch_size) 在內部替你分批、打亂，並對每一批做前向、算 loss、算梯度、更新，Dropout 也會自動在訓練時作用；evaluate 與 predict 則自動用推論模式執行，相當於 PyTorch 的 model.eval() 加 torch.no_grad() 迴圈。Keras 3 可以指定 PyTorch 為運算後端（KERAS_BACKEND="torch"），寫法不變。

下面是一個完整、可以直接執行的 PyTorch 訓練與驗證迴圈，把這一頁的重點都放進去了。輸出是 PyTorch 2.14.1 在 CPU 上的實跑結果（亂數資料，示意）：

import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader

torch.manual_seed(0)
X = torch.randn(500, 8)
y = (X[:, 0] - X[:, 1] > 0).long()                # 類別標籤要是 long
train_dl = DataLoader(TensorDataset(X[:400], y[:400]), batch_size=32, shuffle=True)
val_dl = DataLoader(TensorDataset(X[400:], y[400:]), batch_size=50)

model = nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Dropout(0.2), nn.Linear(32, 2))
loss_fn = nn.CrossEntropyLoss()                   # 內含 softmax，模型輸出 logits
opt = torch.optim.Adam(model.parameters(), lr=0.01)

for epoch in range(5):
    model.train()                                 # 每個 epoch 開頭切回訓練模式
    for xb, yb in train_dl:                       # 400 筆 / 32 → 13 批
        opt.zero_grad()                           # 1. 清梯度
        logits = model(xb)                        # 2. 前向
        loss = loss_fn(logits, yb)                # 3. 算 loss
        loss.backward()                           # 4. 反向，梯度寫進 .grad
        opt.step()                                # 5. 更新參數
    model.eval()                                  # Dropout 關閉
    val_loss, correct = 0.0, 0
    with torch.no_grad():                         # 驗證不建計算圖
        for xb, yb in val_dl:
            logits = model(xb)
            val_loss += loss_fn(logits, yb).item()   # .item() 變成 Python 數字
            correct += (logits.argmax(1) == yb).sum().item()
    print(epoch, round(val_loss / len(val_dl), 4), correct / 100)
# 0 0.4717 0.81 … 4 0.0815 0.98（驗證 loss 逐輪下降）
print(len(train_dl), type(val_loss).__name__)    # 13 float

同樣的事情用 Keras 3 寫（以 PyTorch 為後端實跑），整個雙層迴圈只剩一行 fit：

import os
os.environ["KERAS_BACKEND"] = "torch"            # import keras 之前指定後端
import numpy as np, keras
from keras import layers

rng = np.random.default_rng(0)
X = rng.normal(size=(500, 8)).astype("float32")
y = (X[:, 0] - X[:, 1] > 0).astype("int64")
model = keras.Sequential([keras.Input(shape=(8,)),
                          layers.Dense(32, activation="relu"),
                          layers.Dropout(0.2),          # fit 時作用，evaluate／predict 時自動關閉
                          layers.Dense(2, activation="softmax")])
model.compile(optimizer=keras.optimizers.Adam(0.01),   # 對應 torch.optim.Adam
              loss="sparse_categorical_crossentropy",  # 對應 loss_fn，整數標籤
              metrics=["accuracy"])
hist = model.fit(X[:400], y[:400], epochs=5, batch_size=32, verbose=0)  # 內含整個雙層迴圈
print(len(hist.history["loss"]))                       # 5，每個 epoch 一個 loss
loss, acc = model.evaluate(X[400:], y[400:], verbose=0)  # 對應 eval＋no_grad 的驗證迴圈
print(acc > 0.9)                                       # True

7. zero_grad 可以放在哪裡

容易寫錯或考錯的地方

backward 不會更新權重：loss.backward() 只算梯度並寫進 .grad，真正修改參數的是 optimizer.step()。選項寫「backward 之後權重已經更新」是錯的。

step 在 backward 之前：這一批的梯度還沒算出來，step 用的是 .grad 裡現有的東西：前面剛 zero_grad 過就是 None，參數完全不動；沒清的話就是上一批的梯度，更新永遠慢一拍。程式不會報錯。

梯度裁剪的位置：clip_grad_norm_ 要在 backward 之後、step 之前，放在 step 之後就沒有作用；學習率排程器 scheduler.step() 則通常放在 optimizer.step() 之後。

model.train() 不會訓練：它只切換模式。同樣地，model.eval() 不會關掉梯度，torch.no_grad() 也不會關掉 Dropout，驗證時兩個都要寫；而且下一個 epoch 開頭要記得切回 train。

計數題：一個 epoch 的 step 次數等於 len(train_loader)，也就是樣本數 ÷ batch_size 無條件進位（drop_last=False 時）；總 step 次數再乘上 epoch 數。

loss 與標籤的格式：nn.CrossEntropyLoss 吃的是沒經過 softmax 的 logits，類別標籤要是整數（long）；用 float 的類別標籤會報錯。記錄 loss 用 loss.item()，不要把張量直接加總。

版本：本頁程式以 PyTorch 2.14.1（CPU）與 Keras 3.15 實際執行。PyTorch 2.x 的訓練迴圈寫法與 1.x 相同，主要差異是 2.0 起 zero_grad() 預設把梯度設成 None；torch.compile 等加速功能不改變這五步的順序。

✅ 自我檢測

以下 6 題都是原創的程式閱讀題，輸出都以 PyTorch 2.14.1 實際執行核對。選完會立即顯示對錯與解析，全部作答後會出現總分。目前得分：0 / 6

🎯 重點整理

1. 每個 batch 五步：前向 → 算 loss → backward → step 的順序固定，zero_grad 只要在 backward 之前（或上一輪 step 之後），不能夾在 backward 與 step 中間。
2. backward 只算梯度並「加」進 .grad，step 才修改參數；梯度預設累加，所以每批要清，2.x 的 zero_grad 預設把 .grad 設成 None。
3. 一個 epoch 的 step 次數是 len(train_loader)，即樣本數 ÷ batch_size 無條件進位（drop_last=False）。
4. model.train()／model.eval() 切換 Dropout、BatchNorm 的行為；torch.no_grad() 關掉計算圖。兩件事互不取代，驗證時兩個都要，下一個 epoch 開頭切回 train。
5. 記錄 loss 用 loss.item()；total += loss 會把每一批的計算圖串在一起，記憶體越吃越多。
6. Keras 的 compile 對應建立 optimizer 與 loss_fn，fit 包住整個雙層迴圈與五步，evaluate／predict 自動用推論模式。