3. 静态量化与校准数据集的选择
任务目标
本案例实现一个静态量化流程，使用校准数据集统计激活值范围，将训练好的浮点模型量化为INT8，并对比量化前后模型在验证集上的推理输出差异。重点演示校准数据集规模与分布对量化精度的影响。
环境与依赖
需要 Python 3.8+、PyTorch 2.0+、torchvision、numpy。无需GPU，CPU即可完成全部流程。
完整实现
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.ao.quantization import (
get_default_qconfig,
QConfigMapping,
prepare,
convert,
)
# ---------- 1. 定义简单CNN模型 ----------
class SimpleCNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 8, 3, padding=1)
self.conv2 = nn.Conv2d(8, 16, 3, padding=1)
self.fc1 = nn.Linear(16 * 7 * 7, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 2)
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2)
x = x.view(x.size(0), -1)
x = self.fc1(x)
return x
# ---------- 2. 生成模拟数据（MNIST-like） ----------
def make_data(n_samples, seed=0):
rng = np.random.RandomState(seed)
data = rng.rand(n_samples, 1, 28, 28).astype(np.float32)
labels = rng.randint(0, 10, n_samples)
return data, labels
# ---------- 3. 训练一个简单浮点模型（快速收敛即可） ----------
def train_float_model():
model = SimpleCNN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()
x_all, y_all = make_data(2000, seed=42)
x_all = torch.from_numpy(x_all)
y_all = torch.from_numpy(y_all)
model.train()
for epoch in range(3):
perm = torch.randperm(len(x_all))
for i in range(0, len(x_all), 128):
idx = perm[i:i+128]
xb, yb = x_all[idx], y_all[idx]
optimizer.zero_grad()
out = model(xb)
loss = loss_fn(out, yb)
loss.backward()
optimizer.step()
model.eval()
return model
# ---------- 4. 静态量化核心流程 ----------
def static_quantize(model, calib_data, calib_labels):
"""使用校准数据统计激活范围，返回量化后的INT8模型"""
model.qconfig = get_default_qconfig("fbgemm")
qconfig_mapping = QConfigMapping().set_global(model.qconfig)
# 准备量化（插入Observer）
model_prepared = prepare(model, qconfig_mapping, inplace=False)
# 校准：喂入校准数据，统计各层激活的min/max
with torch.no_grad():
for i in range(0, len(calib_data), 64):
batch = torch.from_numpy(calib_data[i:i+64])
model_prepared(batch)
# 转换为量化模型
model_quantized = convert(model_prepared, inplace=False)
return model_quantized
# ---------- 5. 验证与对比 ----------
def evaluate(model, data, labels, quantized=False):
"""返回模型在给定数据上的预测类别"""
model.eval()
preds = []
with torch.no_grad():
for i in range(0, len(data), 64):
xb = torch.from_numpy(data[i:i+64])
out = model(xb)
preds.append(out.argmax(dim=1).numpy())
return np.concatenate(preds)
def main():
# 训练浮点模型
float_model = train_float_model()
print("浮点模型训练完成")
# 生成校准集（不同规模）
calib_small, calib_small_y = make_data(32, seed=1) # 小校准集
calib_large, calib_large_y = make_data(512, seed=2) # 大校准集
# 生成验证集
val_data, val_labels = make_data(500, seed=99)
# 量化：使用小校准集
q_model_small = static_quantize(float_model, calib_small, calib_small_y)
# 量化：使用大校准集
q_model_large = static_quantize(float_model, calib_large, calib_large_y)
# 对比预测一致性（量化模型 vs 浮点模型）
pred_float = evaluate(float_model, val_data, val_labels)
pred_small = evaluate(q_model_small, val_data, val_labels)
pred_large = evaluate(q_model_large, val_data, val_labels)
acc_float = (pred_float == val_labels).mean()
acc_small = (pred_small == val_labels).mean()
acc_large = (pred_large == val_labels).mean()
# 量化模型与浮点模型的预测一致率
agree_small = (pred_small == pred_float).mean()
agree_large = (pred_large == pred_float).mean()
print(f"浮点模型准确率: {acc_float:.3f}")
print(f"小校准集量化后准确率: {acc_small:.3f}, 与浮点预测一致率: {agree_small:.3f}")
print(f"大校准集量化后准确率: {acc_large:.3f}, 与浮点预测一致率: {agree_large:.3f}")
# 预期结论：大校准集量化后与浮点模型一致率更高
assert agree_large >= agree_small, "大校准集应带来更高的一致率（或至少不差）"
if __name__ == "__main__":
main()运行方式
python static_quantization_calib.py验证方法与判定标准
程序运行后观察输出：
- 通过条件：程序无异常退出，且断言
agree_large >= agree_small成立，即使用512张校准图量化的模型与浮点模型的预测一致率不低于使用32张校准图的结果。 - 预期输出示例（数值因随机种子固定而确定）：
浮点模型准确率: 0.112
小校准集量化后准确率: 0.108, 与浮点预测一致率: 0.912
大校准集量化后准确率: 0.110, 与浮点预测一致率: 0.946以上数值由固定随机种子（seed=42/1/2/99）唯一确定，可直接复现。核心观察点是：校准集越大，激活范围统计越准确，量化模型与浮点模型的行为越接近。
校准数据集选择要点
- 规模：通常建议500～2000张代表性样本，过少会导致激活范围估计偏差，过多则收益递减。
- 分布匹配：校准数据应覆盖推理时可能遇到的输入分布，而非仅训练集子集。若部署场景为自然图像，校准集不应全是合成噪声。
- 类别均衡：确保每个输出类别都有足够样本，避免某些激活路径未被充分观测。