M04-3 | 慢版與近似版：PT2E 的 QAT conv + BN 用的是哪一種

Open In Colab

M04-2 看到 prepare_qat_pt2e 插的子圖。M04 概覽講過，QAT 時 weight 的 fake quant 要量化 \(W_{\text{inf}} = \gamma W/\sigma\)，而讓 BN 仍用 batch 統計有兩種做法，部落格〈Convolution 與 BatchNorm 的融合大法〉整理了兩者：

做法	出處	差別
慢版	量化白皮書（R. Krishnamoorthi, 2018）的修正項做法；PyTorch 的 `_forward_slow`	BN 的 batch 統計取自沒量化的 conv 輸出，所以要多算一次 conv
近似版	PyTorch eager QAT 預設的 `_forward_approximate`（`conv_fused.py`）	batch 統計取自量化後的 conv 輸出，省掉那一次 conv

這本 notebook 把兩種都用 torch.nn.functional 手寫一次，和 PT2E 的圖比：

A. 手寫近似版，和 prepare_qat_pt2e 的圖逐位元比對。
B. 手寫慢版，看它和 PT2E 差多少。
C. 差在哪裡。

import copy

import torch
import torch.nn.functional as F

from common.graph_utils import sqnr
from common.models import ConvBnReLU, example_inputs
from common.pt2e import export_module, get_xnnpack_quantizer, prepare_qat_pt2e, quiet
from common.viz import table

quiet()
_ = torch.manual_seed(0)
EPS = 1e-5  # BatchNorm2d 預設的 ε

建一張 QAT 的 prepared 圖，先讓 BN 的 running stats 與 fake quant 的範圍都有值：

model = ConvBnReLU(3, 8).train()
with torch.no_grad():
    for _ in range(5):
        model(torch.randn(16, 3, 16, 16) * 2 + 1)
quantizer, _ = get_xnnpack_quantizer(is_per_channel=True, is_qat=True)
prepared = prepare_qat_pt2e(export_module(copy.deepcopy(model), example_inputs(), train=True), quantizer)
with torch.no_grad():
    for _ in range(3):
        prepared(torch.randn(16, 3, 16, 16) * 2 + 1)


def state(gm, name):
    """prepared 圖上的 parameter 或 buffer，例如 conv.weight、bn.running_var。"""
    params = dict(gm.named_parameters())
    return params[name] if name in params else gm.get_buffer(name)


def fake_quants(gm):
    """依 forward 順序：輸入 x、折過的 weight、relu 輸出的 fake quant（M04-2 看過這個順序）。"""
    return [gm.get_submodule(n.target) for n in gm.graph.nodes
            if n.op == "call_module" and "activation_post_process" in n.target]

A. 手寫近似版，和 PT2E 的圖比

近似版就是 M04-2 的式 (1)：\(s = \gamma / \sqrt{\sigma^2_{\text{run}} + \varepsilon}\)，先 \(\mathrm{fq}(W \cdot s)\) 做 conv，再除以 \(s\)，最後交給用 batch 統計的 BN。手寫時直接借用 prepared 圖上的三個 fake quant 模組與參數，確保兩邊用的是同一組 scale：

def forward_approximate(gm, x):
    fq_x, fq_w, fq_out = fake_quants(gm)
    w, gamma, beta = state(gm, "conv.weight"), state(gm, "bn.weight"), state(gm, "bn.bias")
    rm, rv = state(gm, "bn.running_mean"), state(gm, "bn.running_var")
    s = gamma / torch.sqrt(rv + EPS)
    y = F.conv2d(fq_x(x), fq_w(w * s.reshape(-1, 1, 1, 1)), None, padding=1) / s.reshape(1, -1, 1, 1)
    y = F.batch_norm(y, rm, rv, gamma, beta, True, 0.1, EPS)  # training=True：用 batch 統計
    return fq_out(F.relu(y))

fake quant 和 BN 在 forward 時都會更新自己的狀態，所以每一條路各用一份獨立的複本：

x = torch.randn(16, 3, 16, 16) * 2 + 1
g_graph, g_approx, g_slow = (copy.deepcopy(prepared) for _ in range(3))
with torch.no_grad():
    y_graph = g_graph(x)
    y_approx = forward_approximate(g_approx, x)
diff = float((y_approx - y_graph).abs().max())
print(f"max|手寫近似版 - prepare_qat_pt2e 的圖| = {diff}")
assert diff == 0

max|手寫近似版 - prepare_qat_pt2e 的圖| = 0.0

逐位元相同：PT2E 的 QAT 子圖就是 PyTorch 的近似版。 torchao 0.18 的原始碼也這樣寫：qat_utils.py 裡 conv + BN 的 pattern 註明是依 nniqat.ConvBn2d._forward_approximate 寫的，而 prepare_qat_pt2e 只有 model、quantizer 兩個參數，沒有切換成慢版的選項。

B. 手寫慢版，看差多少

慢版的 weight fake quant 一樣量化 \(\gamma W / \sigma\)（\(\sigma = \sqrt{\sigma^2_{\text{run}} + \varepsilon}\)，和近似版的 \(W \cdot s\) 相同），差在 batch 統計：先用沒量化的 \(x\)、\(W\) 多算一次 conv，取得這個 batch 的 \(\mu_B\)、\(\sigma_B\)，再把量化的 conv 輸出校正回 batch 統計：

\[ y = \mathrm{conv}\big(\mathrm{fq}(x),\ \mathrm{fq}(\gamma W / \sigma)\big) \cdot \frac{\sigma}{\sigma_B} + \Big(\beta - \frac{\gamma\,\mu_B}{\sigma_B}\Big) \tag{1} \]

def forward_slow(gm, x):
    fq_x, fq_w, fq_out = fake_quants(gm)
    w, gamma, beta = state(gm, "conv.weight"), state(gm, "bn.weight"), state(gm, "bn.bias")
    sigma = torch.sqrt(state(gm, "bn.running_var") + EPS)
    y_float = F.conv2d(x, w, None, padding=1)  # 多出來的那一次 conv：沒量化
    mu_b = y_float.mean((0, 2, 3))
    sigma_b = torch.sqrt(y_float.var((0, 2, 3), unbiased=False) + EPS)
    w_q = fq_w(w * (gamma / sigma).reshape(-1, 1, 1, 1))
    y = F.conv2d(fq_x(x), w_q, None, padding=1) * (sigma / sigma_b).reshape(1, -1, 1, 1)
    y = y + (beta - gamma * mu_b / sigma_b).reshape(1, -1, 1, 1)
    return fq_out(F.relu(y))


with torch.no_grad():
    y_slow = forward_slow(g_slow, x)
    y_fp32 = copy.deepcopy(model).train()(x)  # 不量化、train 模式的原模型
results = {"近似版（= PT2E）": sqnr(y_fp32, y_approx), "慢版": sqnr(y_fp32, y_slow)}
table([[k, f"{v:.2f}"] for k, v in results.items()] + [["慢版 vs 近似版（兩者互比）", f"{sqnr(y_graph, y_slow):.2f}"]],
      ["", "SQNR (dB)，對 fp32 的 train 模式輸出"])

	SQNR (dB)，對 fp32 的 train 模式輸出
近似版（= PT2E）	39.11
慢版	38.89
慢版 vs 近似版（兩者互比）	44.55

gap = abs(results["近似版（= PT2E）"] - results["慢版"])
print(f"兩版對 fp32 的 SQNR 差：{gap:.2f} dB")
assert gap < 1.0

兩版對 fp32 的 SQNR 差：0.22 dB

C. 差在哪裡

fake quant 的輸入完全一樣：慢版的 \(\gamma W / \sigma\) 就是近似版的 \(W \cdot s\)。差別只剩 BN 的 \(\mu_B\)、\(\sigma_B\) 從哪裡來：慢版取自沒量化的 conv，近似版取自量化後的 conv（M04-2 式 (2)）。
所以近似版少算一次 conv；代價是 batch 統計混進了量化誤差（文章的說法）。
這裡的 8-bit、單一 seed 實測，兩版對 fp32 的 SQNR 差不到 1 dB（數字見上表），落在雜訊的量級，不能據此說哪一版比較準。
文章提到在更低的 bit 數（例如 4 bit 以下）差距可能就不能忽略；這裡只測了 8 bit，沒有驗證低 bit 的情況。
eager 的 QAT 模組可以切到慢版（_enable_slow_path_for_better_numerical_stability，預設關閉）；PT2E 沒有這個開關（A 段）。

小結

PT2E 的 QAT conv + BN 子圖和 PyTorch 的 _forward_approximate 逐位元相同，不是白皮書的慢版。
兩版的差別只在 BN 的 batch 統計取自量化前還是量化後；8 bit 下兩者差距很小。