4. 使用ONNX Runtime部署模型

任务目标

将训练好的PyTorch图像分类模型（ResNet18）转换为ONNX格式，并使用ONNX Runtime完成与PyTorch等价的推理，验证两者输出一致性。

环境依赖

torch==2.0.1
torchvision==0.15.2
onnx==1.14.0
onnxruntime==1.15.1
numpy==1.24.3
Pillow==10.0.0

完整实现代码

import torch
import torchvision.models as models
import torch.onnx
import onnx
import onnxruntime as ort
import numpy as np
from PIL import Image

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

# 2. 构造示例输入（1,3,224,224）
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=17,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

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

# 5. 准备测试输入（使用相同随机种子保证可复现）
np.random.seed(42)
test_input = np.random.rand(1, 3, 224, 224).astype(np.float32)

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

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

# 8. 输出一致性验证
max_diff = np.max(np.abs(torch_output - ort_output))
print(f"PyTorch与ONNX Runtime最大输出差异: {max_diff:.6f}")

# 9. 推理耗时对比（预热后各测10次）
def measure_time(func, repeats=10):
    times = []
    for _ in range(repeats):
        start = torch.cuda.Event(enable_timing=True) if torch.cuda.is_available() else None
        import time
        t0 = time.perf_counter()
        func()
        times.append(time.perf_counter() - t0)
    return np.mean(times)

torch_time = measure_time(lambda: model(torch.from_numpy(test_input)))
ort_time = measure_time(lambda: ort_session.run(None, ort_input))
print(f"PyTorch平均推理耗时: {torch_time*1000:.2f} ms")
print(f"ONNX Runtime平均推理耗时: {ort_time*1000:.2f} ms")

# 10. 单张图片推理示例（可选验证）
def preprocess_image(img_path):
    img = Image.open(img_path).resize((224, 224))
    img_array = np.array(img).astype(np.float32) / 255.0
    # 使用ImageNet标准化
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    img_array = (img_array - mean) / std
    img_array = np.transpose(img_array, (2, 0, 1))
    return np.expand_dims(img_array, axis=0).astype(np.float32)

# 加载ImageNet类别标签（仅需前1000行）
import urllib.request
labels_url = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
with urllib.request.urlopen(labels_url) as f:
    labels = [line.strip().decode() for line in f.readlines()]

# 使用示例图片（可用任意图片路径替换）
# img_tensor = preprocess_image("your_image.jpg")
# ort_out = ort_session.run(None, {"input": img_tensor})[0]
# pred_idx = np.argmax(ort_out, axis=1)[0]
# print(f"预测类别: {labels[pred_idx]}")

运行方式

1. 安装依赖：pip install torch torchvision onnx onnxruntime numpy Pillow
2. 直接运行上述Python脚本（保存为deploy_onnx.py后执行python deploy_onnx.py）
3. 脚本将依次完成：模型导出→结构验证→双引擎推理→一致性对比

验证方法与通过标准

• 结构验证：控制台输出“ONNX模型结构验证通过”，说明导出的ONNX文件格式正确
• 数值一致性：PyTorch与ONNX Runtime的最大输出差异应小于1e-4（通常为1e-6量级），证明转换无损
• 动态批处理：导出时设置了动态轴，可用不同batch size（如4、8）的输入测试，ONNX Runtime应能正常推理
• 推理可用性：若使用真实图片测试，输出的预测类别应与PyTorch结果一致

常见问题

• opset版本过低：若导出报错，将opset_version提升至17或更高
• 动态轴不生效：确认dynamic_axes参数已正确设置，且推理时输入shape首维可变化
• CPU/GPU差异：本案例使用CPUExecutionProvider，如需GPU推理将providers改为["CUDAExecutionProvider", "CPUExecutionProvider"]