1. 基于ONNX Runtime的模型转换与推理验证

任务目标

将PyTorch训练好的模型导出为ONNX格式，使用ONNX Runtime在CPU上完成推理，并验证PyTorch与ONNX Runtime的输出一致性。

环境依赖

• Python 3.8+
• torch ≥ 1.12
• onnx ≥ 1.12
• onnxruntime ≥ 1.14
• numpy

完整实现

import torch
import torch.nn as nn
import onnx
import onnxruntime as ort
import numpy as np

# 1. 定义并训练一个简单模型（此处用随机权重模拟训练结果）
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(4, 3)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.fc(x))

model = SimpleModel()
model.eval()

# 2. 固定输入维度，导出ONNX
dummy_input = torch.randn(1, 4)
onnx_path = "simple_model.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes=None,  # 固定shape，便于CPU推理验证
    opset_version=17
)

# 3. 检查ONNX模型结构合法性
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print("ONNX模型结构检查通过")

# 4. ONNX Runtime CPU推理
ort_session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
test_input = torch.randn(1, 4)
ort_input = {ort_session.get_inputs()[0].name: test_input.numpy()}
ort_output = ort_session.run(None, ort_input)[0]

# 5. PyTorch推理（对照）
with torch.no_grad():
    torch_output = model(test_input).numpy()

# 6. 验证输出一致性
np.testing.assert_allclose(torch_output, ort_output, rtol=1e-4, atol=1e-6)
print("PyTorch与ONNX Runtime输出一致，验证通过")
print(f"PyTorch输出: {torch_output}")
print(f"ONNX Runtime输出: {ort_output}")

运行方式

python onnx_convert_verify.py

验证方法与判定标准

• 结构验证：onnx.checker.check_model 无异常，说明导出的ONNX图结构合法。
• 数值一致性：np.testing.assert_allclose 使用 rtol=1e-4、atol=1e-6 对比PyTorch与ONNX Runtime输出，无异常即通过。
• 预期输出：控制台打印“ONNX模型结构检查通过”和“PyTorch与ONNX Runtime输出一致，验证通过”，并显示两组数值相同的输出数组。

常见问题

• 算子不支持：若模型含自定义算子，需在导出时注册算子集或改用 opset_version 适配。
• 动态shape：本案例固定输入shape，若需动态batch，设置 dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}。
• CPU/GPU差异：本案例仅验证CPU推理，若需GPU，将 providers 改为 ["CUDAExecutionProvider"]。