📚 课程目录第4章 ONNX模型导出与静态量化流程

4. ONNX模型导出与静态量化流程

任务说明

本案例完成一个可复现的ONNX模型导出与静态量化闭环：训练一个简易CNN模型用于MNIST手写数字分类，将其导出为ONNX格式，然后使用ONNX Runtime的静态量化API完成INT8量化，并验证量化前后模型的推理输出一致性。

环境依赖

pip install torch torchvision onnx onnxruntime

完整实现代码

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import onnx
import onnxruntime as ort
import numpy as np

# 1. 定义并训练一个简易CNN
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 8, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(8, 16, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.fc = nn.Linear(16 * 7 * 7, 10)

    def forward(self, x):
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        return self.fc(x)

# 数据加载
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_set = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)

# 训练（仅2个epoch用于演示）
model = SimpleCNN()
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
model.train()
for epoch in range(2):
    for images, labels in train_loader:
        optimizer.zero_grad()
        output = model(images)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()
model.eval()

# 2. 导出ONNX模型
dummy_input = torch.randn(1, 1, 28, 28)
onnx_path = "mnist_cnn.onnx"
torch.onnx.export(model, dummy_input, onnx_path,
                  input_names=["input"], output_names=["output"],
                  dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}})

# 3. 静态量化准备：收集校准数据
calib_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True)
calib_data = []
for images, _ in calib_loader:
    calib_data.append(images.numpy())
    if len(calib_data) >= 10:  # 使用320张图像校准
        break

# 4. 执行静态量化
from onnxruntime.quantization import quantize_static, QuantType, CalibrationDataReader

class MNISTDataReader(CalibrationDataReader):
    def __init__(self, data_list):
        self.data_list = data_list
        self.idx = 0
    def get_next(self):
        if self.idx >= len(self.data_list):
            return None
        data = self.data_list[self.idx]
        self.idx += 1
        return {"input": data}

quantized_path = "mnist_cnn_quantized.onnx"
quantize_static(onnx_path, quantized_path, MNISTDataReader(calib_data),
                quant_format=QuantType.QInt8, per_channel=True)

# 5. 验证：比较量化前后输出
def run_inference(model_path, input_data):
    sess = ort.InferenceSession(model_path)
    input_name = sess.get_inputs()[0].name
    return sess.run(None, {input_name: input_data})[0]

test_image = calib_data[0][:1]  # 取一张校准图像
fp32_output = run_inference(onnx_path, test_image)
int8_output = run_inference(quantized_path, test_image)

# 计算预测类别
fp32_pred = np.argmax(fp32_output, axis=1)
int8_pred = np.argmax(int8_output, axis=1)
print(f"FP32预测类别: {fp32_pred[0]}, INT8预测类别: {int8_pred[0]}")
print(f"输出张量形状: {fp32_output.shape} -> {int8_output.shape}")

# 检查模型文件大小
import os
fp32_size = os.path.getsize(onnx_path) / 1024
int8_size = os.path.getsize(quantized_path) / 1024
print(f"FP32模型大小: {fp32_size:.1f} KB, INT8模型大小: {int8_size:.1f} KB")

运行与验证

直接运行上述Python脚本，首次运行会自动下载MNIST数据集（约10MB）
训练过程约需1-2分钟（CPU），训练完成后自动导出ONNX并执行量化
观察输出：应看到FP32与INT8模型对同一张测试图像的预测类别一致（或差异极小），且INT8模型文件明显小于FP32模型

通过条件：脚本无报错执行完毕，输出包含两个模型的预测类别和文件大小信息，且INT8模型文件大小约为FP32模型的1/4左右（静态量化对权重和激活均压缩至8bit）。

关键说明

校准数据使用训练集前320张图像，覆盖了足够的像素分布范围
量化采用per_channel方式，对卷积权重逐通道量化，精度损失更小
动态轴配置允许推理时使用不同batch size，量化模型同样支持
若预测类别不一致，可增加训练epoch或增加校准数据量（将calib_data上限从10提升至20）