1. 模型量化基础

任务目标

使用ONNX Runtime的静态量化API，将一个预训练的FP32图像分类模型（ResNet18）转换为INT8量化模型，并在Tiny ImageNet验证集上对比量化前后Top-1准确率，确认精度损失在可接受范围内（<2%）。

环境与依赖

  • Python 3.9+
  • PyTorch 2.0+（用于导出ONNX模型）
  • onnx 1.14+
  • onnxruntime 1.16+
  • numpy、Pillow

实现步骤

本案例采用ONNX Runtime的静态量化方案，核心流程为：FP32模型导出 → 校准数据准备 → 静态量化 → 精度验证。静态量化需要少量校准数据统计激活值范围，因此我们使用Tiny ImageNet验证集的前200张图片作为校准集，后800张作为测试集。

关键完整代码

import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
import onnx
import onnxruntime as ort
from onnxruntime.quantization import quantize_static, QuantType, CalibrationMethod
from onnxruntime.quantization import CalibrationDataReader
import os
import glob

# ---------- 1. 导出FP32 ONNX模型 ----------
def export_fp32_onnx():
    model = models.resnet18(pretrained=True)
    model.eval()
    dummy_input = torch.randn(1, 3, 224, 224)
    torch.onnx.export(model, dummy_input, "resnet18_fp32.onnx",
                      input_names=["input"], output_names=["output"],
                      opset_version=13)
    print("FP32 ONNX模型已导出")

# ---------- 2. 准备校准与测试数据 ----------
class TinyImageNetReader(CalibrationDataReader):
    def __init__(self, img_dir, batch_size=8):
        self.img_paths = sorted(glob.glob(os.path.join(img_dir, "*.JPEG")))[:200]
        self.batch_size = batch_size
        self.idx = 0
        self.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])
        ])
        self.input_name = "input"

    def get_next(self):
        if self.idx >= len(self.img_paths):
            return None
        batch = []
        for _ in range(self.batch_size):
            if self.idx >= len(self.img_paths):
                break
            img = Image.open(self.img_paths[self.idx]).convert("RGB")
            batch.append(self.transform(img).numpy())
            self.idx += 1
        return {self.input_name: np.stack(batch).astype(np.float32)}

# ---------- 3. 执行静态量化 ----------
def quantize_model():
    calib_reader = TinyImageNetReader("tiny_imagenet_val")
    quantize_static(
        model_input="resnet18_fp32.onnx",
        model_output="resnet18_int8.onnx",
        calibration_data_reader=calib_reader,
        quant_format=QuantType.QOperator,
        per_channel=True,
        activation_type=QuantType.QInt8,
        weight_type=QuantType.QInt8,
        calibrate_method=CalibrationMethod.MinMax
    )
    print("INT8量化模型已生成")

# ---------- 4. 精度验证 ----------
def evaluate_accuracy(model_path, img_dir, num_samples=800):
    sess = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
    input_name = sess.get_inputs()[0].name
    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])
    ])
    img_paths = sorted(glob.glob(os.path.join(img_dir, "*.JPEG")))[200:200+num_samples]
    correct = 0
    for path in img_paths:
        img = Image.open(path).convert("RGB")
        tensor = transform(img).unsqueeze(0).numpy()
        pred = sess.run(None, {input_name: tensor})[0]
        label = int(path.split("_")[-1].split(".")[0])
        if np.argmax(pred) == label:
            correct += 1
    return correct / len(img_paths)

if __name__ == "__main__":
    # 数据目录需包含tiny_imagenet_val/*.JPEG，文件名格式: xxx_标签ID.JPEG
    export_fp32_onnx()
    quantize_model()
    fp32_acc = evaluate_accuracy("resnet18_fp32.onnx", "tiny_imagenet_val")
    int8_acc = evaluate_accuracy("resnet18_int8.onnx", "tiny_imagenet_val")
    print(f"FP32 Top-1准确率: {fp32_acc:.4f}")
    print(f"INT8 Top-1准确率: {int8_acc:.4f}")
    print(f"精度差: {fp32_acc - int8_acc:.4f}")
    assert fp32_acc - int8_acc < 0.02, "精度损失超过2%，量化失败"
    print("量化验证通过")

运行方式

  1. 下载Tiny ImageNet验证集（20000张图片），按文件名格式类别_标签ID.JPEG存放于tiny_imagenet_val/目录，其中标签ID为0-199的整数。
  2. 安装依赖：pip install torch torchvision onnx onnxruntime pillow numpy
  3. 执行脚本：python quantize_resnet18.py

验证方法

脚本运行结束后会输出FP32与INT8模型的Top-1准确率及差值。当且仅当精度差<2%时，程序断言通过并打印"量化验证通过"。同时可检查生成的两个文件：resnet18_fp32.onnx（约45MB）与resnet18_int8.onnx（约12MB），INT8模型体积应约为FP32的1/4。

预期输出

FP32 ONNX模型已导出
INT8量化模型已生成
FP32 Top-1准确率: 0.5234
INT8 Top-1准确率: 0.5178
精度差: 0.0056
量化验证通过

注：以上准确率为示例值，实际结果取决于校准数据分布与测试集，但精度差通常小于2%。

常见问题

  • 校准数据不足：若校准集少于100张，可能导致量化后精度骤降，建议使用200张以上。
  • 预处理不一致：校准与测试时的图像预处理必须与模型训练时一致（Resize→CenterCrop→Normalize），否则精度异常。
  • 算子不支持：若模型包含量化不支持的算子（如某些自定义OP），可尝试quant_format=QuantType.QDQ格式。
← 上一章 📖 返回目录 下一章 →