2. PyTorch模型动态量化实现与精度评估

任务目标

本案例将带领你使用 PyTorch 对训练好的 CNN 模型（以 ResNet18 为例）实施动态量化（Dynamic Quantization），将模型权重从 FP32 压缩为 INT8，并完成量化前后模型的精度与体积对比评估。动态量化无需校准数据集，适合 LSTM、Transformer 等层以及全连接层占主导的模型，是嵌入式部署中最快速的量化入门路径。

环境依赖

需要 torch、torchvision、numpy。若需 GPU 加速训练基线模型，请安装对应 CUDA 版本；量化与评估过程支持纯 CPU 完成。

关键实现与代码

以下代码完整实现了：构建基线模型 → 加载预训练权重 → 定义动态量化配置 → 执行量化 → 在测试集上对比精度 → 对比模型文件体积。核心量化 API 为 torch.quantization.quantize_dynamic，仅需一行调用即可完成。

import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from torchvision.datasets import CIFAR10
from torch.utils.data import DataLoader
import os, time, copy

# ---------- 1. 模型定义与预处理 ----------
def get_model():
    model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
    # 将最后的全连接层替换为 CIFAR-10 的 10 分类
    model.fc = nn.Linear(model.fc.in_features, 10)
    return model

transform_test = transforms.Compose([
    transforms.Resize(224),          # ResNet18 需要 224x224 输入
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# ---------- 2. 加载测试数据（仅用 2000 张子集加速评估） ----------
testset = CIFAR10(root='./data', train=False, download=True, transform=transform_test)
test_subset = torch.utils.data.Subset(testset, range(2000))
testloader = DataLoader(test_subset, batch_size=64, shuffle=False, num_workers=2)

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

# ---------- 4. 基线模型（FP32） ----------
model_fp32 = get_model()
# 此处仅演示量化流程，实际部署时应加载在 CIFAR-10 上微调后的权重
# 若直接使用 ImageNet 预训练权重，精度会偏低，但量化对比趋势仍然有效

# ---------- 5. 执行动态量化 ----------
model_int8 = copy.deepcopy(model_fp32)
model_int8.eval()

# 关键 API：动态量化，仅量化权重为 int8，激活仍为浮点
quantized_model = torch.quantization.quantize_dynamic(
    model_int8,                    # 原始模型
    {nn.Linear, nn.Conv2d},        # 量化层类型（ResNet18 的卷积层也可被量化）
    dtype=torch.qint8              # 量化精度
)

# ---------- 6. 对比评估 ----------
print("=== 精度对比（CIFAR-10 测试子集 2000 张）===")
fp32_acc = evaluate(model_fp32, testloader)
int8_acc = evaluate(quantized_model, testloader)
print(f"FP32 模型精度: {fp32_acc:.2f}%")
print(f"INT8 动态量化模型精度: {int8_acc:.2f}%")
print(f"精度损失: {fp32_acc - int8_acc:.2f} 个百分点")

# ---------- 7. 模型体积对比 ----------
os.makedirs('./quant_models', exist_ok=True)
fp32_path = './quant_models/resnet18_fp32.pth'
int8_path = './quant_models/resnet18_int8_dynamic.pth'
torch.save(model_fp32.state_dict(), fp32_path)
torch.save(quantized_model.state_dict(), int8_path)

fp32_size = os.path.getsize(fp32_path) / 1024 / 1024
int8_size = os.path.getsize(int8_path) / 1024 / 1024
print(f"\n=== 模型体积对比 ===")
print(f"FP32 模型大小: {fp32_size:.2f} MB")
print(f"INT8 动态量化模型大小: {int8_size:.2f} MB")
print(f"体积压缩比: {fp32_size / int8_size:.2f}x")

# ---------- 8. 推理速度对比（CPU 上运行 200 batch） ----------
def measure_inference_time(model, loader, num_batches=50):
    model.eval()
    start = time.time()
    with torch.no_grad():
        for i, (images, _) in enumerate(loader):
            if i >= num_batches:
                break
            _ = model(images)
    return time.time() - start

print("\n=== CPU 推理耗时对比（50 batch）===")
t_fp32 = measure_inference_time(model_fp32, testloader)
t_int8 = measure_inference_time(quantized_model, testloader)
print(f"FP32 推理耗时: {t_fp32:.3f}s")
print(f"INT8 动态量化推理耗时: {t_int8:.3f}s")
print(f"推理加速比: {t_fp32 / t_int8:.2f}x")

运行与验证方法

直接运行上述脚本即可。首次运行会自动下载 CIFAR-10 数据集（约 170MB）和 ResNet18 预训练权重。由于未在 CIFAR-10 上微调，基线精度约 10%~20%（接近随机猜测），但量化前后的相对精度损失趋势和体积压缩比仍然有效。

通过判定标准

• 量化执行成功：脚本无报错，输出 INT8 模型精度与 FP32 精度。
• 体积压缩：INT8 模型文件体积约为 FP32 的 1/3 至 1/4（压缩比 3~4x）。
• 精度损失：动态量化通常带来 1%~3% 的精度下降（若使用微调后的权重，损失会更小）。若精度下降超过 5%，说明模型对量化敏感，需考虑量化感知训练。
• 推理加速：在 CPU 上 INT8 模型推理耗时通常比 FP32 快 1.5~3 倍（取决于硬件对 INT8 指令的支持程度）。

常见问题排查

• 精度过低：确认是否加载了在目标数据集上微调过的权重。使用 ImageNet 预训练权重直接测试 CIFAR-10 属于跨域测试，精度本身就不高，但量化对比趋势仍可参考。
• 量化后推理变慢：检查 CPU 是否支持 AVX2/AVX512 指令集（PyTorch 的 INT8 内核依赖这些指令）。在老旧 CPU 上可能无加速效果。
• 内存不足：减少测试子集数量（如改为 500 张）或减小 batch_size。