📚 课程目录第2章静态量化与校准

2. 静态量化与校准

任务目标

实现一个最小可复现的静态量化流程：使用PyTorch对预训练模型进行静态量化（per-tensor、对称、int8），用代表性校准数据集确定激活值的量化范围，最后对比量化前后模型在验证集上的精度与权重/激活张量的量化参数。

环境依赖

Python 3.8+，PyTorch ≥ 1.13（torch.ao.quantization API）
torchvision（用于获取预训练模型与数据集）

实现与关键代码

以下代码完整实现了：加载预训练MobileNetV2 → 替换为量化可感知架构 → 用CIFAR-10训练集子集做校准（统计激活范围）→ 执行静态量化 → 在测试集上对比精度。核心量化配置为MinMaxObserver（权重与激活均使用），量化粒度per-tensor，dtype为qint8/quint8。

import torch
import torchvision
import torchvision.transforms as T
from torch.ao.quantization import (
    get_default_qconfig, prepare_qat, convert, QConfigMapping
)
from torch.ao.quantization.observer import MinMaxObserver

# ---------- 1. 数据准备（CIFAR-10，归一化到[0,1]） ----------
transform = T.Compose([T.ToTensor()])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
calib_loader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(testset, batch_size=256, shuffle=False)

# ---------- 2. 加载预训练模型并融合 ----------
model = torchvision.models.mobilenet_v2(weights=torchvision.models.MobileNet_V2_Weights.IMAGENET1K_V1)
model.eval()
# 融合Conv+BN+ReLU（静态量化必需）
model.fuse_model()

# ---------- 3. 配置静态量化 ----------
qconfig = QConfigMapping().set_global(
    torch.ao.quantization.QConfig(
        activation=MinMaxObserver.with_args(dtype=torch.quint8, qscheme=torch.per_tensor_affine),
        weight=MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric)
    )
)
model.qconfig = qconfig
# 插入观察者（校准模式）
model_prepared = prepare_qat(model, qconfig)  # 复用QAT的prepare以保留BN统计
model_prepared.eval()

# ---------- 4. 校准：用代表性数据统计激活范围 ----------
print("Calibrating...")
with torch.no_grad():
    for i, (images, _) in enumerate(calib_loader):
        model_prepared(images)
        if i >= 20:  # 使用约1280张图片校准
            break

# ---------- 5. 执行静态量化转换 ----------
model_quantized = convert(model_prepared)

# ---------- 6. 精度对比 ----------
def evaluate(model, loader):
    correct, total = 0, 0
    with torch.no_grad():
        for images, labels in loader:
            out = model(images)
            pred = out.argmax(dim=1)
            correct += (pred == labels).sum().item()
            total += labels.size(0)
    return correct / total

acc_fp32 = evaluate(model, test_loader)
acc_int8 = evaluate(model_quantized, test_loader)
print(f"FP32 accuracy: {acc_fp32:.4f}")
print(f"INT8 static-quantized accuracy: {acc_int8:.4f}")

# ---------- 7. 查看量化参数（验证校准生效） ----------
# 取第一个量化卷积层
for name, mod in model_quantized.named_modules():
    if hasattr(mod, 'weight_fake_quant') or hasattr(mod, 'activation_post_process'):
        if hasattr(mod, 'weight_fake_quant'):
            wq = mod.weight_fake_quant
            print(f"Layer {name}: weight scale={wq.scale.item():.6f}, zero_point={wq.zero_point.item()}")
        if hasattr(mod, 'activation_post_process'):
            aq = mod.activation_post_process
            print(f"Layer {name}: act scale={aq.scale.item():.6f}, zero_point={aq.zero_point.item()}")
        break  # 仅演示第一层

运行方式

pip install torch torchvision
python static_quantization.py

验证方法与判定标准

校准生效验证：程序输出的第一个量化层权重scale应为非零小值（约0.001～0.01量级），激活scale约为0.01～0.1量级。若scale全为0或未打印，说明校准未正确执行。
精度对比：INT8精度应接近FP32（通常下降<2%）。由于CIFAR-10与ImageNet预训练分布存在差异，FP32基线本身可能只有60%～70%左右，重点观察量化前后差值。
确定性输出参考：使用上述固定校准数据（前20个batch）时，量化参数是确定性的。示例输出（非实测）：FP32 accuracy: 0.6821，INT8 accuracy: 0.6715，实际数值取决于PyTorch版本与随机数据加载顺序。

常见问题

校准batch数影响：校准数据过少（<200张）会导致激活范围估计不足，精度下降明显；过多则增加校准时间。1280张（20 batch×64）是精度与速度的平衡点。
模型不支持的层：若遇到不支持的算子（如某些自定义层），需在量化前将其替换为可量化版本或保持浮点运行（通过.qconfig=None排除）。
BN层处理：静态量化要求先融合BN，否则量化后精度会显著下降。本代码已通过fuse_model()完成。