3. 图像分类模型INT8量化与精度评估
任务目标
对预训练图像分类模型(ResNet18)执行INT8量化,使用端侧推理引擎(以ONNX Runtime + XNNPACK EP模拟端侧环境)运行,对比量化前后Top-1精度,验证精度损失在可接受范围内(<1%)。
环境与依赖
- Python 3.8+,PyTorch 1.13+(用于加载预训练模型与校准数据)
- onnxruntime 1.15+(含XNNPACK EP,支持INT8量化模型推理)
- onnx、onnxruntime-tools(用于模型转换与量化)
- torchvision(提供ResNet18预训练权重与ImageNet数据预处理)
实现方案
采用静态量化流程:先加载PyTorch预训练ResNet18并转换为ONNX FP32模型,再使用校准数据集统计激活值范围完成INT8量化。端侧推理通过ONNX Runtime的XNNPACK执行提供程序运行量化模型,该EP针对ARM/x86端侧CPU优化,可代表典型端侧推理环境。
关键完整代码
import torch
import torchvision.models as models
import onnx
import onnxruntime as ort
from onnxruntime.quantization import quantize_static, QuantType
from onnxruntime.quantization.calibrate import CalibrationDataReader
import numpy as np
from torchvision import datasets, transforms
# ---------- 1. 导出FP32 ONNX模型 ----------
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
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"],
dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}})
# ---------- 2. 校准数据读取器(使用ImageNet验证集子集) ----------
class CalibReader(CalibrationDataReader):
def __init__(self, batch_size=16, num_batches=10):
self.batch_size = batch_size
self.num_batches = num_batches
self.batch_idx = 0
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])
])
# 使用ImageNet验证集(需提前下载至./imagenet_val)
self.dataset = datasets.ImageFolder("./imagenet_val", transform=transform)
self.loader = torch.utils.data.DataLoader(
self.dataset, batch_size=batch_size, shuffle=True)
def get_next(self):
if self.batch_idx >= self.num_batches:
return None
batch = next(iter(self.loader))
self.batch_idx += 1
return {"input": batch[0].numpy()}
# ---------- 3. 执行静态INT8量化 ----------
calib_reader = CalibReader()
quantize_static(
model_input="resnet18_fp32.onnx",
model_output="resnet18_int8.onnx",
calibration_data_reader=calib_reader,
quant_format=QuantType.QInt8,
per_channel=True,
weight_type=QuantType.QInt8,
activation_type=QuantType.QInt8
)
# ---------- 4. 精度评估(FP32 vs INT8) ----------
def evaluate_accuracy(model_path, use_xnnpack=True):
providers = ["XnnpackExecutionProvider", "CPUExecutionProvider"] if use_xnnpack \
else ["CPUExecutionProvider"]
sess = ort.InferenceSession(model_path, providers=providers)
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])
])
dataset = datasets.ImageFolder("./imagenet_val", transform=transform)
loader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=False)
correct = 0
total = 0
with torch.no_grad():
for images, labels in loader:
outputs = sess.run(None, {"input": images.numpy()})[0]
preds = np.argmax(outputs, axis=1)
correct += np.sum(preds == labels.numpy())
total += labels.size(0)
return correct / total
fp32_acc = evaluate_accuracy("resnet18_fp32.onnx", use_xnnpack=False)
int8_acc = evaluate_accuracy("resnet18_int8.onnx", use_xnnpack=True)
print(f"FP32 Top-1 Accuracy: {fp32_acc:.4f}")
print(f"INT8 Top-1 Accuracy: {int8_acc:.4f}")
print(f"Accuracy Drop: {fp32_acc - int8_acc:.4f}")
运行与验证
- 准备数据:下载ImageNet验证集(至少160张图片,按类别目录存放),路径为
./imagenet_val。 - 执行脚本:
python quantize_resnet18.py - 通过条件:控制台输出INT8模型Top-1精度与FP32差值小于0.01(即精度损失<1%),且INT8模型推理正常无报错。
预期输出示例
FP32 Top-1 Accuracy: 0.6981
INT8 Top-1 Accuracy: 0.6944
Accuracy Drop: 0.0037
(注:以上数值仅为基于公开ResNet18 ImageNet精度的合理示例,实际结果取决于校准数据与验证子集,以实际运行为准。)
常见问题
- XNNPACK不支持某些算子:若报错,可回退至CPUExecutionProvider验证功能正确性,但精度评估需在XNNPACK下完成以代表端侧环境。
- 校准数据不足:建议至少使用100张以上覆盖不同类别的图片,否则激活值范围统计不准会导致精度显著下降。
- 动态轴影响量化:若需固定batch=1可移除dynamic_axes,量化过程更稳定。