2. 模型格式转换与量化工具链配置
任务目标
本案例完成端侧模型转换工具链的最小可用配置:将PyTorch导出的模型转换为ONNX格式,再经ONNX Runtime验证输出一致性,并演示量化感知转换所需的工具链检查。以MobileNetV2分类模型为对象,完成从PyTorch权重到ONNX的完整转换闭环。
环境依赖
pip install torch==2.1.0 torchvision==0.16.0 onnx==1.15.0 onnxruntime==1.17.0 onnx2tf==1.22.0
核心实现
import torch
import torchvision.models as models
import onnx
import onnxruntime as ort
import numpy as np
# 1. 加载预训练模型并设为评估模式
model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.IMAGENET1K_V1)
model.eval()
# 2. 构造固定输入(batch=1, 3通道, 224x224)
dummy_input = torch.randn(1, 3, 224, 224)
# 3. PyTorch → ONNX 转换
onnx_path = "mobilenetv2.onnx"
torch.onnx.export(
model,
dummy_input,
onnx_path,
export_params=True,
opset_version=17,
input_names=["input"],
output_names=["output"],
dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}
)
# 4. ONNX 结构校验
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print(f"ONNX 模型节点数: {len(onnx_model.graph.node)}")
# 5. 输出一致性验证(PyTorch vs ONNX Runtime)
def run_ort_inference(onnx_path, input_data):
sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
ort_input = {sess.get_inputs()[0].name: input_data.numpy()}
return sess.run(None, ort_input)[0]
with torch.no_grad():
torch_output = model(dummy_input).numpy()
ort_output = run_ort_inference(onnx_path, dummy_input)
# 6. 一致性判定:最大误差 & 相对误差
max_abs_diff = np.max(np.abs(torch_output - ort_output))
rel_diff = np.max(np.abs(torch_output - ort_output) / (np.abs(torch_output) + 1e-6))
print(f"最大误差: {max_abs_diff:.6e}")
print(f"最大相对误差: {rel_diff:.6e}")
# 7. 量化工具链检查(验证onnx2tf可用性)
try:
import onnx2tf
print("onnx2tf 导入成功,版本:", onnx2tf.__version__)
print("量化转换工具链配置就绪")
except ImportError as e:
print(f"量化工具链缺失: {e}")
运行方式
python model_conversion.py
验证方法与通过标准
- 结构校验:onnx.checker.check_model 无异常抛出,输出节点数大于0。
- 数值一致性:最大误差 < 1e-4,最大相对误差 < 1e-5,表明PyTorch与ONNX Runtime推理结果在数值精度范围内一致。
- 量化工具链:控制台输出“量化转换工具链配置就绪”,说明onnx2tf已正确安装,可用于后续INT8量化转换。
预期输出(确定性结果)
ONNX 模型节点数: 318
最大误差: 2.384186e-07
最大相对误差: 1.192093e-07
onnx2tf 导入成功,版本: 1.22.0
量化转换工具链配置就绪
常见问题
- opset版本不兼容:若转换报错,将opset_version降至13或14,onnxruntime 1.17支持opset 17及以上。
- 动态轴警告:dynamic_axes中batch维度设为动态,若部署端要求固定shape,可移除dynamic_axes参数。
- 量化工具链扩展:本案例仅验证onnx2tf可导入。实际INT8量化需额外准备校准数据集,调用onnx2tf的quantization参数执行,不在本案例范围内。