📚 课程目录第2章基于ONNX Runtime的模型转换与部署

2. 基于ONNX Runtime的模型转换与部署

任务目标

将训练好的PyTorch模型（以ResNet18图像分类为例）导出为ONNX格式，并在车载Linux设备（x86_64或ARM64）上使用ONNX Runtime完成推理，验证转换前后输出一致性。

环境与依赖

PyTorch ≥ 1.12（含torchvision）
onnx ≥ 1.13
onnxruntime ≥ 1.14（车载端建议使用对应架构的wheel包）
NumPy、Pillow（用于图像预处理）

实现步骤与代码

以下代码完成三个核心动作：构造并导出ONNX模型、使用ONNX Runtime加载推理、对比PyTorch原始输出。导出时固定batch维度为1，并显式指定动态轴以适配车载端可变输入尺寸（若需固定尺寸可移除dynamic_axes参数）。

import torch
import torchvision.models as models
import onnx
import onnxruntime as ort
import numpy as np

# ---------- 1. 构造并导出ONNX模型 ----------
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)
onnx_path = "resnet18.onnx"

torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}
)

# ---------- 2. 校验ONNX模型结构 ----------
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print(f"ONNX模型结构校验通过，节点数: {len(onnx_model.graph.node)}")

# ---------- 3. ONNX Runtime推理与一致性对比 ----------
ort_session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])

# 使用相同输入进行推理
test_input = torch.randn(1, 3, 224, 224)

# PyTorch原始输出
with torch.no_grad():
    torch_output = model(test_input).numpy()

# ONNX Runtime输出
ort_input = {ort_session.get_inputs()[0].name: test_input.numpy()}
ort_output = ort_session.run(None, ort_input)[0]

# 对比最大误差
max_diff = np.max(np.abs(torch_output - ort_output))
print(f"PyTorch与ONNX Runtime输出最大误差: {max_diff:.6f}")

# 判定一致性（误差小于1e-4视为通过）
assert max_diff < 1e-4, "转换前后输出差异过大"
print("一致性验证通过：ONNX模型输出与PyTorch原始模型一致")

运行方式

在具备PyTorch的开发机上运行上述脚本，生成resnet18.onnx文件。
将resnet18.onnx拷贝至车载设备（ARM64或x86_64）。
在车载设备上安装onnxruntime：pip install onnxruntime（ARM64请使用对应架构wheel）。
车载端仅需执行上述代码中第3步（ONNX Runtime加载与推理），无需PyTorch环境。

验证方法与判定标准

结构验证：onnx.checker.check_model无异常，说明ONNX图结构合法。
数值一致性：PyTorch与ONNX Runtime对同一随机输入的最大误差小于1e-4，说明转换未引入精度损失。
车载端可用性：在车载设备上能成功创建InferenceSession并完成一次前向推理，即证明部署链路打通。

常见问题

算子不支持：若导出时报错，可降低opset_version至13或11重试。
动态尺寸：车载端若需固定输入尺寸，可删除dynamic_axes参数，导出后输入shape固定为(1,3,224,224)。
推理速度：ONNX Runtime默认使用CPUExecutionProvider；若车载端有NPU/GPU，可改用对应执行提供程序（如CUDAExecutionProvider或TensorrtExecutionProvider）。