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参数执行,不在本案例范围内。