import copy
import torch
from common.data import calibration_batches
from common.graph_utils import observer_nodes, qdq_nodes, short, sqnr
from common.models import example_inputs, get_model
from common.pt2e import convert_pt2e, export_module, get_xnnpack_quantizer, prepare_pt2e, quiet
from common.viz import draw, table
quiet()
_ = torch.manual_seed(0)M01-3 | QuantizationConfig 的三個開關
M01-1、M01-2 的 quantizer 都用 get_xnnpack_quantizer(is_per_channel=True) 建立, 底下是 get_symmetric_quantization_config(...) 產生一份 QuantizationConfig,再 set_global 給整個模型。 這份 config 有三個開關:
| 參數 | 意思 |
|---|---|
is_per_channel |
weight 每個 output channel 各一組 scale(False:整個 weight 共用一組) |
is_qat |
observer 換成 fake quant,forward 時真的做 quantize → dequantize,給 QAT 訓練用 |
is_dynamic |
activation 不在校準時決定 qparams,改成執行時依當下的輸入計算 |
這本 notebook 用同一個模型、同一份校準資料跑過各種組合,看圖和數值差在哪。
A. 四種設定並排
run_one() 把一個設定從 export 跑到 convert,回傳:插了幾個 observer、型別、q/dq 數量、輸出對 fp32 的 SQNR。 模型是沒訓練過的 TinyConvNet,這裡只比較量化前後的輸出,不看 accuracy。
def run_one(model, ex, **cfg):
quantizer, _ = get_xnnpack_quantizer(**cfg)
prepared = prepare_pt2e(export_module(copy.deepcopy(model), ex), quantizer)
obs_types = sorted({type(m).__name__ for _, m in observer_nodes(prepared)})
n_obs = len(observer_nodes(prepared))
with torch.no_grad():
for batch in calibration_batches(n=256, batch_size=32):
prepared(*batch)
converted = convert_pt2e(prepared)
with torch.no_grad():
out_sqnr = sqnr(model(*ex), converted(*ex))
return [n_obs, ", ".join(obs_types), len(qdq_nodes(converted)), f"{out_sqnr:.1f}"]
model = get_model("tiny").eval()
ex = example_inputs(8)
configs = {
"per-tensor weight": dict(is_per_channel=False),
"per-channel weight": dict(is_per_channel=True),
"per-channel + is_qat": dict(is_per_channel=True, is_qat=True),
"per-channel + is_dynamic": dict(is_per_channel=True, is_dynamic=True),
}
results = {name: run_one(model, ex, **cfg) for name, cfg in configs.items()}
table([[name, *r] for name, r in results.items()],
["設定", "#observer", "observer 型別", "#q/dq", "輸出 SQNR (dB)"])| 設定 | #observer | observer 型別 | #q/dq | 輸出 SQNR (dB) |
|---|---|---|---|---|
| per-tensor weight | 16 | HistogramObserver, MinMaxObserver | 28 | 20.6 |
| per-channel weight | 16 | HistogramObserver, PerChannelMinMaxObserver | 28 | 20.6 |
| per-channel + is_qat | 16 | FusedMovingAvgObsFakeQuantize | 28 | 17.4 |
| per-channel + is_dynamic | 2 | PerChannelMinMaxObserver, PlaceholderObserver | 3 | 50.3 |
三個值得追問的地方:
- per-tensor 和 per-channel 的 SQNR 幾乎一樣。per-channel 沒用嗎?(B 段)
is_dynamic只剩 2 個 observer,SQNR 卻高很多。(C 段)is_qat把 observer 換成FusedMovingAvgObsFakeQuantize。它是給訓練用的,M06 才會真的拿來訓練; 這裡只要記得「同一個 quantizer,換一個 config,插進圖裡的東西就不同」。
B. per-tensor ≈ per-channel:誤差主要來自 activation
只量化 weight、activation 保持 fp32,看輸出差多少。 做法是直接把模型裡每個 conv/linear 的 weight 換成 quantize → dequantize 之後的值 (BN 還沒折進 conv,只是近似 PT2E 的做法)。
def fake_quant_weight(w, per_channel):
"""對稱 int8:scale = absmax / 127,per-channel 時每個 output channel 各算一個。"""
dims = tuple(range(1, w.dim()))
scale = (w.abs().amax(dim=dims, keepdim=True) if per_channel else w.abs().max()) / 127
return (w / scale).round().clamp(-127, 127) * scale
def weight_only_sqnr(model, ex, per_channel):
m = copy.deepcopy(model)
with torch.no_grad():
for mod in m.modules():
if isinstance(mod, (torch.nn.Conv2d, torch.nn.Linear)):
mod.weight.copy_(fake_quant_weight(mod.weight, per_channel))
return sqnr(model(*ex), m(*ex))
rows = [[name, f"{weight_only_sqnr(model, ex, pc):.1f}", results[full][3]]
for name, pc, full in [("per-tensor", False, "per-tensor weight"),
("per-channel", True, "per-channel weight")]]
table(rows, ["weight 粒度", "只量化 weight 的 SQNR (dB)", "完整量化的 SQNR (dB)"])| weight 粒度 | 只量化 weight 的 SQNR (dB) | 完整量化的 SQNR (dB) |
|---|---|---|
| per-tensor | 48.7 | 20.6 |
| per-channel | 46.1 | 20.6 |
for _, w_only, full in rows:
assert float(w_only) > float(full) + 20只量化 weight 時,兩種粒度的 SQNR 都比完整量化高出 20 dB 以上:這個模型的誤差主要來自 activation, 所以 weight 用哪種粒度在最後的 SQNR 上看不出差別。 另一個原因是這個模型剛初始化,每個 channel 的 weight 範圍差不多。
per-channel 的價值要在「channel 之間範圍差很多」時才看得出來。造一個 weight,讓第 0 個 channel 大 100 倍:
w = torch.randn(4, 3, 3, 3)
w[0] *= 100.0
w_pt, w_pc = fake_quant_weight(w, per_channel=False), fake_quant_weight(w, per_channel=True)
rows = [[c, f"{float(w[c].abs().max()):.2f}", f"{sqnr(w[c], w_pt[c]):.1f}", f"{sqnr(w[c], w_pc[c]):.1f}"]
for c in range(w.shape[0])]
rows.append(["整個 tensor", "", f"{sqnr(w, w_pt):.1f}", f"{sqnr(w, w_pc):.1f}"])
table(rows, ["channel", "absmax", "per-tensor SQNR (dB)", "per-channel SQNR (dB)"])| channel | absmax | per-tensor SQNR (dB) | per-channel SQNR (dB) |
|---|---|---|---|
| 0 | 302.52 | 44.0 | 44.0 |
| 1 | 3.20 | 4.2 | 44.4 |
| 2 | 2.55 | 4.9 | 45.7 |
| 3 | 2.23 | 2.4 | 45.6 |
| 整個 tensor | 38.5 | 44.0 |
per-tensor 只有一個 scale,由最大的 channel 0 決定(absmax / 127)。 其他 channel 的數值都不到一格 scale 的幾倍,量化後幾乎全被壓成 0,SQNR 只剩個位數。 「整個 tensor」那一個 row的 SQNR 被 channel 0 主導,看不出這個問題;要逐 channel 看。
weight 的 per-channel 幾乎沒有代價:scale 是常數,跟著 weight 存。activation 的 per-channel 則大多數硬體不支援。
C. is_dynamic:只量化了 linear
dynamic 的 SQNR 高,是因為量化的東西少。先看 annotate 的結果:
gm = export_module(copy.deepcopy(model), ex)
q_dyn, _ = get_xnnpack_quantizer(is_per_channel=True, is_dynamic=True)
annotated = q_dyn.annotate(copy.deepcopy(gm))
marked = [n.name for n in annotated.graph.nodes if "quantization_annotation" in n.meta]
print("有標註的 node:", marked)
assert "linear" in marked and not any(name.startswith("conv2d") for name in marked)有標註的 node: ['fc_weight', 'fc_bias', 'linear']
只有 linear(和它的 weight、bias)有標註,conv 一個都沒有。原因寫在 quantizer 的原始碼裡 (torchao 0.18 的 torchao/testing/pt2e/_xnnpack_quantizer.py):
annotate()一看到 global config 是 dynamic,就整個改走動態的標註路徑;conv、add 這些靜態的 annotator 完全不會執行。- 動態路徑只處理
DYNAMIC_OPS = ["linear"]這個清單。
所以這裡的 dynamic 不是「linear 用 static、conv 改成 dynamic」,而是「只有 linear 被量化,而且用動態的方式」。 跑完 convert,看圖上剩下哪些運算:
prepared_dyn = prepare_pt2e(gm, q_dyn)
with torch.no_grad():
prepared_dyn(*ex)
converted_dyn = convert_pt2e(prepared_dyn)
print(" → ".join(short(n.target) for n in converted_dyn.graph.nodes if n.op == "call_function"))dequantize_per_channel → conv2d → relu → max_pool2d → conv2d → relu → conv2d → add → relu → max_pool2d → conv2d → relu → adaptive_avg_pool2d → flatten → choose_qparams → getitem → getitem → quantize_per_tensor → dequantize_per_tensor → linear
convs = [n for n in converted_dyn.graph.nodes if short(n.target) == "conv2d"]
assert all(a.op != "call_function" or "quantize" not in str(a.target) for c in convs for a in c.all_input_nodes)
assert any(short(n.target) == "choose_qparams" for n in converted_dyn.graph.nodes)
print(f"{len(convs)} 個 conv2d 的輸入都沒有 dequantize:conv 用 fp32 跑")4 個 conv2d 的輸入都沒有 dequantize:conv 用 fp32 跑
conv2d 前面沒有任何 dq,用 fp32 的 weight 直接算。量化只發生在 linear 附近:
draw(converted_dyn, focus=["linear"], hops=4, title="is_dynamic=True 的 convert 圖:linear 附近")同一個 linear 上,weight 和輸入用的是兩種不同的方式:
| 對象 | 量化方式 |
|---|---|
| linear 的 weight | 靜態:convert 時就存成 int8 常數(per-channel),只剩 dequantize_per_channel |
| linear 的輸入 activation | 動態:prepare 時插的是 PlaceholderObserver(不收集統計,只佔位);convert 後變成 choose_qparams → quantize_per_tensor → dequantize_per_tensor,每次執行時由 choose_qparams 依當下的輸入算 scale/zp |
| conv 等其他 op | 不量化,沒有 q/dq |
也就是說,dynamic 指的是「activation 的 qparams 在執行時才算、不需要校準」,不是把所有 op 換成另一種量化方式; 哪些 op 支援動態量化,由 quantizer 決定。換一個 quantizer,答案就可能不同: executorch 附的 XNNPACKQuantizer(M08 會用到)支援動態量化的 pattern 有 linear、conv、conv_transpose (讀 executorch 1.5 原始碼;conv_relu 不在其中)。
回到 A 段的表:dynamic 那一個 row 只量化了 linear,其他三個 row量化了整個模型,所以 SQNR 不能直接比。
小結
- config 決定的是「插什麼 observer、量化哪些 op」,同一個 quantizer 換 config,圖就不同。
- 比較 SQNR 前,先確認量化了哪些 op,否則比較的對象不同。
- per-channel weight 在 channel 範圍差異大時才關鍵;要逐 channel 看誤差,整體數字會被最大的 channel 蓋過。