2. 基于PTQ的8位激活量化

任务目标

对已训练好的浮点模型实施PTQ（训练后量化）中的8位激活量化：使用少量校准数据统计激活值分布，据此确定量化范围（scale和zero_point），将激活从FP32映射到INT8，并对比量化前后模型输出的数值差异，验证量化误差在可接受范围内。

环境与依赖

本案例使用PyTorch 2.0+和torchvision，无需额外安装。量化采用PyTorch官方torch.quantization API实现，校准数据使用CIFAR-10训练集前128个batch（每个batch 64张图）。

完整实现

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.quantization import prepare_qat, convert
from torch.quantization import get_default_qconfig, QConfig
from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx
from torch.ao.quantization import MinMaxObserver, PerChannelMinMaxObserver
from torch.ao.quantization.observer import HistogramObserver
import copy

# 1. 定义一个小型CNN（与CIFAR-10匹配）
class SmallCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 8 * 8, 64),
            nn.ReLU(),
            nn.Linear(64, 10),
        )

    def forward(self, x):
        return self.classifier(self.features(x))

# 2. 准备校准数据（CIFAR-10验证集前128个batch）
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
calib_dataset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform)
calib_loader = torch.utils.data.DataLoader(
    calib_dataset, batch_size=64, shuffle=False,
    sampler=torch.utils.data.SubsetRandomSampler(range(128 * 64)))

# 3. 加载预训练模型（此处用随机权重演示流程，实际应加载训练好的权重）
model_fp32 = SmallCNN()
model_fp32.eval()

# 4. 配置量化：MinMaxObserver确定激活范围
qconfig = QConfig(
    activation=MinMaxObserver.with_args(dtype=torch.quint8, qscheme=torch.per_tensor_affine),
    weight=PerChannelMinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_channel_symmetric)
)

# 5. 使用FX图模式量化（支持更灵活的量化范围统计）
from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx
model_to_quantize = copy.deepcopy(model_fp32)
qconfig_dict = {"": qconfig}
model_prepared = prepare_fx(model_to_quantize, qconfig_dict)

# 6. 校准：统计激活值min/max
print("开始校准（统计激活量化范围）...")
with torch.no_grad():
    for i, (images, _) in enumerate(calib_loader):
        model_prepared(images)
        if i % 32 == 0:
            print(f"  校准进度: {i}/{len(calib_loader)}")
print("校准完成")

# 7. 转换为量化模型
model_quantized = convert_fx(model_prepared)
print("量化模型转换完成")

# 8. 验证：比较量化前后输出差异
def get_outputs(model, loader, num_batches=5):
    """返回前num_batches个batch的模型输出（已排序，便于比较）"""
    outputs = []
    with torch.no_grad():
        for i, (images, _) in enumerate(loader):
            if i >= num_batches:
                break
            outputs.append(model(images))
    return torch.cat(outputs, dim=0)

# 使用独立的验证集（非校准数据）
val_loader = torch.utils.data.DataLoader(
    torchvision.datasets.CIFAR10(
        root='./data', train=False, download=True, transform=transform),
    batch_size=64, shuffle=False)

print("计算量化前输出...")
out_fp32 = get_outputs(model_fp32, val_loader, num_batches=5)
print("计算量化后输出...")
out_int8 = get_outputs(model_quantized, val_loader, num_batches=5)

# 9. 计算差异指标
abs_diff = (out_fp32 - out_int8).abs()
cos_sim = nn.functional.cosine_similarity(out_fp32, out_int8, dim=1).mean()

print("\n===== 量化前后输出差异 =====")
print(f"最大误差: {abs_diff.max().item():.6f}")
print(f"平均误差: {abs_diff.mean().item():.6f}")
print(f"余弦相似度:   {cos_sim.item():.6f}")

# 10. 查看量化参数（激活的scale和zero_point）
print("\n===== 激活量化参数（第1个卷积层后） =====")
for name, module in model_quantized.named_modules():
    if hasattr(module, 'activation_post_process'):
        obs = module.activation_post_process
        print(f"  {name}: scale={obs.scale.item():.6f}, zero_point={obs.zero_point.item()}")
        break  # 只显示第一个

运行方式

  1. 直接运行脚本：python ptq_activation_quant.py（首次运行会自动下载CIFAR-10数据集，约170MB）
  2. 若已有训练好的模型权重，将model_fp32替换为加载权重的模型即可

验证方法与判定标准

运行后观察终端输出，满足以下条件即判定量化有效：

  • 余弦相似度 > 0.99：表明量化前后输出方向基本一致，信息未严重丢失
  • 平均误差 < 0.1：在CIFAR-10的10类logits输出上，0.1以内的平均偏差属于正常PTQ误差范围
  • 激活scale值合理：打印出的scale应在1e-3 ~ 1e-1量级（取决于激活值范围），zero_point通常接近128（对于ReLU后的非负激活）

若余弦相似度低于0.95，说明量化范围统计不准确，可尝试：增加校准batch数、改用HistogramObserver（百分位截断）替代MinMaxObserver，或对模型做BN折叠。

关键说明

  • 本案例使用FX图模式量化（prepare_fx/convert_fx），相比旧版torch.quantization.prepare能自动处理ReLU等算子的融合，且校准统计更精确
  • 校准数据仅用于统计激活值的min/max，不参与梯度更新，因此PTQ无需训练
  • 示例代码使用随机权重模型仅用于演示量化流程；实际使用时请替换为训练好的权重，量化误差会更有参考意义
  • 若需部署到实际硬件，可进一步导出为ONNX或使用torch.jit脚本化模型
← 上一章 📖 返回目录 下一章 →