M06-2 | 評估 QAT 模型前要切兩個開關：BN 的 train／eval，與 observer

Open In Colab

M05-2 看過：PTQ 的圖沒有 BN，校準不受 batch 組成影響。QAT 的圖不一樣：

• BN 還在（M04-2），train 模式下用的是當下 batch 的統計。
• fake quant 裡的 observer 只要開著，每次 forward 都會更新 min／max，scale 跟著變。

這兩件事都讓「同一筆樣本」的輸出，隨著「和誰一起跑」「之前跑過什麼」而改變。

• A. 三種狀態下，同一筆樣本單獨跑和同批跑差多少，同一份測試集評估兩次的 loss 是否相同。
• B. allow_exported_model_train_eval：把 .train()／.eval() 接回 GraphModule。

import copy

import torch

from common.data import calibration_batches, evaluate_loss, make_loader, trained_tiny_net
from common.models import example_inputs
from common.pt2e import (
    allow_exported_model_train_eval,
    disable_observer,
    enable_observer,
    export_module,
    move_exported_model_to_eval,
    move_exported_model_to_train,
    prepare_qat_pt2e,
    quiet,
)
from common.qat import qat_quantizer, train_qat
from common.viz import table

quiet()
_ = torch.manual_seed(0)

先訓練一個 int8 的 QAT 模型。train_qat() 是本專案的工具（common/qat.py），內容就是 M06-1 B 段逐行寫的迴圈，多了一個每個 step 都會呼叫的 hook；結束時模型處於 eval 模式、observer 仍開著。

model = trained_tiny_net()
test_loader = make_loader(n=1024, seed=2, shuffle=False)
(x16,) = next(calibration_batches(16, seed=5))

prepared = prepare_qat_pt2e(export_module(model, example_inputs(), train=True), qat_quantizer(bits=8))
_ = train_qat(prepared, epochs=3)

A. 三種狀態

兩個開關：

• BN：move_exported_model_to_train／move_exported_model_to_eval（M06-1）。
• observer：model.apply(enable_observer)／model.apply(disable_observer)，切的是每個 fake quant 的 observer_enabled buffer（M02-2）。另一個開關 fake_quant_enabled 這裡不動。

每種狀態量兩件事：

• 單獨 vs 同批：16 筆樣本逐筆單獨跑，和 16 筆一起跑，輸出的最大差。
• 評估兩次：同一份測試集連續評估兩次，test loss 是否相同。

每次都在 deepcopy 上跑，不讓前一項的副作用影響下一項。

@torch.no_grad()
def alone_vs_together(gm, x):
    together = copy.deepcopy(gm)(x)
    g = copy.deepcopy(gm)
    alone = torch.cat([g(x[i : i + 1]) for i in range(len(x))])
    return float((alone - together).abs().max())


STATES = {
    "train 模式、observer 開": (move_exported_model_to_train, enable_observer),
    "eval 模式、observer 開": (move_exported_model_to_eval, enable_observer),
    "eval 模式、observer 關": (move_exported_model_to_eval, disable_observer),
}
results = {}
for name, (set_bn, set_observer) in STATES.items():
    g = copy.deepcopy(prepared)
    set_bn(g)
    g.apply(set_observer)
    diff = alone_vs_together(g, x16)
    results[name] = (diff, evaluate_loss(g, test_loader), evaluate_loss(g, test_loader))
table([[name, f"{d:.1e}", f"{l1:.5f}", f"{l2:.5f}", "是" if l1 == l2 else "否"] for name, (d, l1, l2) in results.items()],
      ["狀態", "單獨 vs 同批：最大差", "test loss 第 1 次", "第 2 次", "可重現"])

狀態	單獨 vs 同批：最大差	test loss 第 1 次	第 2 次	可重現
train 模式、observer 開	6.1e+00	0.87850	0.88116	否
eval 模式、observer 開	3.0e-01	0.83002	0.82649	否
eval 模式、observer 關	0.0e+00	0.82809	0.82809	是

reproducible = [l1 == l2 for _, l1, l2 in results.values()]
assert reproducible == [False, False, True]
assert results["train 模式、observer 開"][0] > results["eval 模式、observer 開"][0] > 0

• train 模式：BN 用 batch 統計，同一筆樣本的輸出隨同批的其他樣本改變。評估本身還會更新 running stats，所以連評估兩次都不一樣。
• eval 模式、observer 開：BN 固定了，但每次 forward 都在更新 fake quant 的 min／max（moving average），scale 跟著動，輸出也跟著動。評估資料偷偷參與了校準。
• eval 模式、observer 關：兩次評估完全一致，這才是可以拿來報告的數字。

單獨 vs 同批在第三種狀態偶爾仍會差「恰好一格」：CPU 上 batch 1 與 batch 16 的 conv 走不同的實作，約 1e-5 的浮點差，剛好把落在 round 邊界的值推到隔壁格。這和量化流程無關，見 M07-1。

所以評估或 convert 之前固定做兩件事：move_exported_model_to_eval 與 apply(disable_observer)。 本專案把這兩行包成 freeze_for_eval()（common/qat.py），之後的 notebook 都用它。訓練中途要評估也一樣，評完再 enable_observer、切回 train。

B. allow_exported_model_train_eval

不想每次都寫 move_exported_model_to_train，可以用 allow_exported_model_train_eval(gm) 把 GraphModule 的 .train()／.eval() 換成這兩個函式。下面讀圖上 4 個 batch_norm node 的 training 參數（第 6 個位置參數）：

def bn_training_flags(gm):
    return [n.args[5] for n in gm.graph.nodes if n.op == "call_function" and "batch_norm" in str(n.target)]


g = copy.deepcopy(prepared)
flags = {"（訓練後）": bn_training_flags(g)}
allow_exported_model_train_eval(g)
g.train()
flags[".train()"] = bn_training_flags(g)
g.eval()
flags[".eval()"] = bn_training_flags(g)
table([[k, str(v)] for k, v in flags.items()], ["呼叫", "batch_norm 的 training 參數"])

呼叫	batch_norm 的 training 參數
（訓練後）	[False, False, False, False]
.train()	[True, True, True, True]
.eval()	[False, False, False, False]

assert [set(v) for v in flags.values()] == [{False}, {True}, {False}]

• 它只是把 .train()／.eval() 換成 move_exported_model_to_train／eval（讀 torchao 0.18 原始碼），影響的仍然只有 BN 與 dropout，observer 不受影響。
• model.training 這個屬性和一般 nn.Module 的語意不同，不要拿它判斷狀態。
• 好處是「為一般 nn.Module 寫的」訓練／評估函式可以直接用在 QAT 圖上。本章都明寫 move_exported_model_to_train／eval，讓切換的時機看得見。

小結

• 評估或 convert 前：BN 切 eval、observer 關掉（freeze_for_eval）。少做一個，結果就不可重現。
• allow_exported_model_train_eval 讓 .train()／.eval() 能用，但仍不管 observer。