📚 课程目录第2章 PyTorch模型静态量化实战

2. PyTorch模型静态量化实战

任务目标

对预训练 ResNet-18 执行 PyTorch 官方静态量化流程（per-tensor、8-bit、对称量化），使用 ImageNet 验证集子集完成校准与精度评估，对比 FP32 与 INT8 模型的 Top-1 准确率，并输出量化后模型文件。

环境依赖

Python 3.8+，PyTorch ≥1.13（含 torchvision），tqdm。无需额外硬件（CPU 即可完成校准与推理）。

实现代码

import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Subset
from torchvision.datasets import ImageFolder
import os, random, tqdm

# ---------- 1. 配置 ----------
DATA_ROOT = "/path/to/imagenet_subset"  # 子集目录：train/ 和 val/ 各含类别子文件夹
BATCH_SIZE = 32
NUM_CALIB_BATCHES = 20          # 校准用 batch 数
NUM_EVAL_SAMPLES = 2000         # 评估样本数（从 val 中随机抽取）
SEED = 42
random.seed(SEED)
torch.manual_seed(SEED)

# ---------- 2. 数据预处理 ----------
calib_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])
])

# 校准集：使用训练子集（或单独的校准集）
calib_dataset = ImageFolder(os.path.join(DATA_ROOT, "train"), transform=calib_transform)
calib_loader = DataLoader(calib_dataset, batch_size=BATCH_SIZE, shuffle=True)

# 评估集：验证子集
eval_dataset = ImageFolder(os.path.join(DATA_ROOT, "val"), transform=calib_transform)
indices = random.sample(range(len(eval_dataset)), min(NUM_EVAL_SAMPLES, len(eval_dataset)))
eval_subset = Subset(eval_dataset, indices)
eval_loader = DataLoader(eval_subset, batch_size=BATCH_SIZE, shuffle=False)

# ---------- 3. 定义量化配置 ----------
def get_qconfig():
    return torch.quantization.get_default_qconfig('fbgemm')  # x86 平台

# ---------- 4. 模型准备与静态量化 ----------
def prepare_model_for_quantization():
    model = models.resnet18(pretrained=True)
    model.eval()
    # 融合 Conv+BN+ReLU
    model = torch.quantization.fuse_modules(model, [['conv1', 'bn1', 'relu'],
                                                     ['layer1.0.conv1', 'layer1.0.bn1', 'layer1.0.relu'],
                                                     ['layer1.0.conv2', 'layer1.0.bn2'],
                                                     ['layer2.0.conv1', 'layer2.0.bn1', 'layer2.0.relu'],
                                                     ['layer2.0.conv2', 'layer2.0.bn2'],
                                                     ['layer3.0.conv1', 'layer3.0.bn1', 'layer3.0.relu'],
                                                     ['layer3.0.conv2', 'layer3.0.bn2'],
                                                     ['layer4.0.conv1', 'layer4.0.bn1', 'layer4.0.relu'],
                                                     ['layer4.0.conv2', 'layer4.0.bn2']])
    model.qconfig = get_qconfig()
    torch.quantization.prepare(model, inplace=True)
    return model

def calibrate(model, loader, num_batches):
    print("校准中...")
    with torch.no_grad():
        for i, (inputs, _) in enumerate(loader):
            if i >= num_batches:
                break
            model(inputs)

def evaluate(model, loader):
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in tqdm.tqdm(loader, desc="评估中"):
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return 100.0 * correct / total

# ---------- 5. 执行量化 ----------
# FP32 基线
fp32_model = models.resnet18(pretrained=True).eval()
fp32_acc = evaluate(fp32_model, eval_loader)
print(f"FP32 Top-1 准确率: {fp32_acc:.2f}%")

# 静态量化
q_model = prepare_model_for_quantization()
calibrate(q_model, calib_loader, NUM_CALIB_BATCHES)
torch.quantization.convert(q_model, inplace=True)

# 量化后评估
int8_acc = evaluate(q_model, eval_loader)
print(f"INT8 静态量化 Top-1 准确率: {int8_acc:.2f}%")
print(f"精度损失: {fp32_acc - int8_acc:.2f} 个百分点")

# 保存量化模型
torch.save(q_model.state_dict(), "resnet18_static_quantized.pth")
print("量化模型已保存至 resnet18_static_quantized.pth")

运行方式

准备 ImageNet 子集目录结构：train/ 与 val/ 下各含类别子文件夹，每类至少 20 张以上图片；
修改 DATA_ROOT 为实际路径；
执行脚本：python static_quantization.py。

验证方法与通过条件

校准过程无报错，校准后模型成功执行 convert；
FP32 与 INT8 模型均在相同 2000 张验证子集上完成评估，输出准确率数值；
INT8 Top-1 与 FP32 相比损失通常在 0.5~2 个百分点以内（具体取决于子集分布），若损失超过 5 个百分点需检查校准集与评估集分布是否一致；
生成的 resnet18_static_quantized.pth 可被 torch.load 正常加载。

常见问题

fbgemm 后端不可用：若在 ARM 平台运行，将 get_qconfig() 改为 torch.quantization.get_default_qconfig('qnnpack')；
融合层不完整：ResNet-18 的 BasicBlock 中第二个卷积后无 ReLU（恒等映射相加前），因此仅融合 conv2+bn2，不融合 relu；
校准 batch 数影响：NUM_CALIB_BATCHES 过少（<10）可能导致量化参数不准，建议 20~50。