2. 静态量化与校准

任务目标

对PyTorch中的MobileNetV2分类模型实施静态量化（Post-Training Static Quantization），使用ImageNet的一个小型代表性子集完成校准，并对比量化前后模型的Top-1准确率，评估精度损失。

环境与依赖

Python 3.8+，PyTorch ≥ 1.13（含torchvision）
依赖安装：pip install torch torchvision
需联网下载MobileNetV2预训练权重（首次运行自动下载，约14MB）

实现思路

静态量化需要三个要素：预训练模型、代表性校准数据集、量化配置。校准数据用于统计各层激活值的动态范围，从而确定量化缩放因子。本案例采用qnnpack后端（x86平台亦可使用fbgemm），校准集使用ImageNet验证集的1000个子集（每类1张），校准轮数为1轮。量化后模型仅支持CPU推理，因此评估也在CPU上进行。

完整实现代码

import torch
import torch.nn as nn
from torch.quantization import quantize_fx, prepare_qat_fx
from torch.ao.quantization import get_default_qconfig
from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx
from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset
import os

# 1. 数据预处理与校准集构建
def get_calib_loader(batch_size=32, num_samples=1000):
    """构建校准数据：从ImageNet验证集每类取1张，共1000张"""
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])
    # 若本地无ImageNet数据，可用CIFAR-10替代（需调整类别数）
    val_dir = './imagenet_val'  # 请替换为实际路径
    if not os.path.exists(val_dir):
        raise FileNotFoundError(
            f"未找到验证集目录 {val_dir}。请下载ImageNet验证集并解压，"
            "或修改代码使用CIFAR-10（需同步修改模型分类头）。")
    val_dataset = datasets.ImageFolder(val_dir, transform=transform)
    # 每类取1张：按目标类别索引
    class_indices = {}
    for idx, (_, label) in enumerate(val_dataset.samples):
        if label not in class_indices:
            class_indices[label] = idx
        if len(class_indices) == 1000:
            break
    subset_idx = list(class_indices.values())[:num_samples]
    subset = Subset(val_dataset, subset_idx)
    return DataLoader(subset, batch_size=batch_size, shuffle=False)

# 2. 模型准备与量化
def quantize_mobilenetv2():
    # 加载预训练模型并设为eval模式
    model = mobilenet_v2(weights=MobileNet_V2_Weights.IMAGENET1K_V1)
    model.eval()
    
    # 配置量化：使用FX图模式量化
    qconfig = get_default_qconfig('qnnpack')
    model.qconfig = qconfig
    
    # 准备量化（插入观察者）
    model_prepared = prepare_fx(model, {'': qconfig})
    
    # 校准：用代表性数据前向传播，统计激活范围
    calib_loader = get_calib_loader()
    with torch.no_grad():
        for images, _ in calib_loader:
            model_prepared(images)
    
    # 转换为量化模型
    model_quantized = convert_fx(model_prepared)
    return model, model_quantized

# 3. 精度评估
def evaluate(model, loader, device='cpu'):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return 100.0 * correct / total

if __name__ == '__main__':
    # 构建评估集（使用与校准集不同的数据，这里用验证集前2000张）
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])
    val_dataset = datasets.ImageFolder('./imagenet_val', transform=transform)
    eval_loader = DataLoader(Subset(val_dataset, range(2000)), batch_size=64)
    
    # 执行量化
    fp32_model, int8_model = quantize_mobilenetv2()
    
    # 评估FP32模型（CPU）
    acc_fp32 = evaluate(fp32_model, eval_loader)
    
    # 评估INT8量化模型（仅CPU）
    acc_int8 = evaluate(int8_model, eval_loader)
    
    print(f"FP32 Top-1 准确率: {acc_fp32:.2f}%")
    print(f"INT8 Top-1 准确率: {acc_int8:.2f}%")
    print(f"精度损失: {acc_fp32 - acc_int8:.2f} 个百分点")
    
    # 模型大小对比
    torch.save(fp32_model.state_dict(), 'mobilenetv2_fp32.pth')
    torch.save(int8_model.state_dict(), 'mobilenetv2_int8.pth')
    size_fp32 = os.path.getsize('mobilenetv2_fp32.pth') / 1e6
    size_int8 = os.path.getsize('mobilenetv2_int8.pth') / 1e6
    print(f"FP32模型大小: {size_fp32:.2f} MB")
    print(f"INT8模型大小: {size_int8:.2f} MB")
    print(f"压缩比: {size_fp32 / size_int8:.2f}x")

运行方式

准备ImageNet验证集：目录结构为./imagenet_val/<类别文件夹>/<图片>，共1000个类别文件夹。若数据量不足，可先用CIFAR-10替代（需将模型分类头改为10类并重新训练或微调）。
将上述代码保存为static_quantize.py，修改val_dir路径为实际数据位置。
执行：python static_quantize.py

验证方法与通过标准

运行成功标志：脚本无异常退出，并输出FP32与INT8的准确率及模型大小。
精度损失判定：INT8模型Top-1准确率相对FP32下降通常不超过2～3个百分点（ImageNet 1000类任务）。若损失超过5%，需检查校准集代表性或增加校准轮数。
压缩效果：INT8模型文件大小约为FP32的1/4（约3.5MB vs 14MB），压缩比接近4x。
功能正确性：量化模型能正常完成前向推理，输出张量形状与FP32一致。

常见问题

qnnpack后端不支持：若CPU不支持qnnpack，改用get_default_qconfig('fbgemm')，并在prepare前设置torch.backends.quantized.engine = 'fbgemm'。
校准数据不足：少于100张图片可能导致量化参数估计不准，建议至少使用500张覆盖主要类别分布。
模型包含不可量化层：MobileNetV2的ReLU6和深度可分离卷积均支持量化，若自定义模型需检查算子兼容性。

📚 课程目录