1. PyTorch模型量化基础

任务目标

使用PyTorch对预训练ResNet-18执行动态量化（Dynamic Quantization），将模型权重从FP32压缩为INT8，并对比量化前后的模型文件大小与CPU推理速度。

环境依赖

  • Python 3.8+
  • PyTorch ≥ 1.8（含torchvision）
  • CPU环境即可（动态量化主要面向CPU推理优化）

核心实现

import torch
import torchvision.models as models
import time
import copy

# 1. 加载预训练模型并设为评估模式
model_fp32 = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
model_fp32.eval()

# 2. 深拷贝一份用于量化
model_quantized = copy.deepcopy(model_fp32)

# 3. 动态量化：仅量化Linear层（ResNet-18的FC层）
model_quantized = torch.ao.quantization.quantize_dynamic(
    model_quantized, {torch.nn.Linear}, dtype=torch.qint8
)

# 4. 模型大小对比（保存后查看文件字节数）
import os
def get_model_size_mb(model, filename):
    torch.save(model.state_dict(), filename)
    size_bytes = os.path.getsize(filename)
    os.remove(filename)
    return size_bytes / (1024 * 1024)

size_fp32 = get_model_size_mb(model_fp32, "resnet18_fp32.pth")
size_int8 = get_model_size_mb(model_quantized, "resnet18_int8.pth")
print(f"FP32模型大小: {size_fp32:.2f} MB")
print(f"INT8动态量化模型大小: {size_int8:.2f} MB")
print(f"压缩比: {size_fp32 / size_int8:.2f}x")

# 5. CPU推理速度对比（使用随机输入）
batch_size = 1
dummy_input = torch.randn(batch_size, 3, 224, 224)

def measure_inference_time(model, input_tensor, num_runs=50):
    # 预热
    with torch.no_grad():
        for _ in range(10):
            model(input_tensor)
    # 计时
    torch.cuda.synchronize() if torch.cuda.is_available() else None
    start = time.time()
    with torch.no_grad():
        for _ in range(num_runs):
            model(input_tensor)
    end = time.time()
    return (end - start) / num_runs * 1000  # 毫秒

# 限制线程数保证公平对比
torch.set_num_threads(4)

time_fp32 = measure_inference_time(model_fp32, dummy_input)
time_int8 = measure_inference_time(model_quantized, dummy_input)
print(f"FP32平均推理耗时: {time_fp32:.2f} ms")
print(f"INT8动态量化平均推理耗时: {time_int8:.2f} ms")
print(f"加速比: {time_fp32 / time_int8:.2f}x")

# 6. 验证输出一致性（量化前后预测差异）
with torch.no_grad():
    out_fp32 = model_fp32(dummy_input)
    out_int8 = model_quantized(dummy_input)
# 比较最大误差
max_diff = (out_fp32 - out_int8).abs().max().item()
print(f"量化前后输出最大误差: {max_diff:.6f}")
print("验证完成：量化模型可正常推理且输出与FP32模型接近。")

运行方式

# 安装依赖
pip install torch torchvision

# 运行脚本
python quantize_resnet18.py

验证方法与判定标准

  • 模型大小：量化后模型文件应显著小于FP32模型（通常压缩比约3~4倍，因仅量化全连接层，ResNet-18卷积层仍为FP32）。
  • 推理速度：量化后模型在CPU上的推理耗时应当低于或接近FP32模型（动态量化主要优化全连接层计算，若整体耗时下降不明显，可观察FC层占比高的输入场景）。
  • 输出一致性：量化前后模型对同一输入的输出最大误差应小于0.1（典型值在1e-3~1e-2量级），说明量化精度损失在可接受范围。
  • 功能正确性：脚本无报错，且打印出上述对比指标即视为运行成功。

关键说明

  • 动态量化仅对torch.nn.Linear层生效，ResNet-18的卷积层仍保持FP32，因此压缩比和加速比有限。若需更大压缩，可改用静态量化（需校准数据）或量化感知训练。
  • 推理速度受CPU型号、线程数影响，建议固定torch.set_num_threads()以保证对比公平。
  • 若需部署到移动端或边缘设备，可进一步导出为ONNX并量化，但本案例聚焦PyTorch原生动态量化流程。
← 上一章 📖 返回目录 下一章 →