1. PyTorch模型转换为ONNX格式

任务说明

本案例将一个训练好的PyTorch图像分类模型（ResNet18）转换为ONNX格式，并通过对比转换前后模型在相同输入下的输出，验证转换一致性。转换后的ONNX文件将作为后续量化、推理部署的基础。

环境依赖

需要安装：torch、torchvision、onnx、onnxruntime。建议使用Python 3.8+。

完整实现代码

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

# 1. 加载预训练模型并设为评估模式
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
model.eval()

# 2. 构造示例输入（1个batch，3通道，224x224）
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出为ONNX格式
onnx_path = "resnet18.onnx"
torch.onnx.export(
    model,                      # 待转换模型
    dummy_input,                # 示例输入
    onnx_path,                  # 输出文件路径
    export_params=True,         # 保存模型权重
    opset_version=13,           # ONNX算子集版本
    input_names=["input"],      # 输入节点名
    output_names=["output"],    # 输出节点名
    dynamic_axes={              # 动态batch维度
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)
print(f"ONNX模型已保存至: {onnx_path}")

# 4. 验证ONNX模型结构完整性
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print("ONNX模型结构检查通过")

# 5. 使用ONNX Runtime进行推理，对比输出一致性
# 准备相同输入
test_input = torch.randn(1, 3, 224, 224)

# PyTorch模型推理
with torch.no_grad():
    torch_output = model(test_input).numpy()

# ONNX Runtime推理
ort_session = ort.InferenceSession(onnx_path)
ort_input = {ort_session.get_inputs()[0].name: test_input.numpy()}
ort_output = ort_session.run(None, ort_input)[0]

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

# 7. 验证通过条件：误差小于1e-4
assert max_diff < 1e-4, "输出一致性验证失败"
print("输出一致性验证通过")

运行方式

将上述代码保存为 convert_onnx.py，在终端执行：

python convert_onnx.py

验证方法与判定标准

• 结构验证：程序执行到“ONNX模型结构检查通过”且无异常，说明ONNX文件格式正确。
• 一致性验证：程序输出“输出一致性验证通过”，且最大误差小于1e-4，即认为转换成功。
• 预期输出示例（数值因随机输入而异）：

ONNX模型已保存至: resnet18.onnx
ONNX模型结构检查通过
PyTorch与ONNX输出最大误差: 2.384186e-07
输出一致性验证通过

常见问题

• 算子不支持：若导出时报算子不支持错误，可尝试降低opset_version（如改为11）。
• 动态维度失效：若后续推理需要固定batch，可移除dynamic_axes参数。
• 误差过大：检查模型是否处于eval()模式，且输入预处理方式与训练时一致。