📚 课程目录第3章使用ONNX Runtime加速模型推理

3. 使用ONNX Runtime加速模型推理

任务说明

本案例完成一个可复现的推理加速闭环：训练一个简单的PyTorch分类模型 → 导出为ONNX格式 → 使用ONNX Runtime加载并推理 → 对比PyTorch与ONNX Runtime的推理速度和精度差异。精度对比以相同输入下两者输出是否一致为判定标准，速度对比以多次推理的平均耗时为参考。

环境依赖

pip install torch onnx onnxruntime numpy

完整实现代码

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

# 1. 定义并训练一个简单模型（两分类MLP）
class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(4, 16),
            nn.ReLU(),
            nn.Linear(16, 2)
        )
    def forward(self, x):
        return self.fc(x)

torch.manual_seed(0)
model = SimpleMLP()
model.eval()

# 用随机数据模拟训练（实际使用时替换为真实训练过程）
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()
X_train = torch.randn(200, 4)
y_train = torch.randint(0, 2, (200,))
for epoch in range(50):
    optimizer.zero_grad()
    out = model(X_train)
    loss = loss_fn(out, y_train)
    loss.backward()
    optimizer.step()

# 2. 导出ONNX
dummy_input = torch.randn(1, 4)
onnx_path = "simple_mlp.onnx"
torch.onnx.export(model, dummy_input, onnx_path,
                  input_names=["input"], output_names=["output"],
                  dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}})

# 验证ONNX模型结构
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print("ONNX模型导出并检查通过")

# 3. ONNX Runtime推理
ort_session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

# 4. 精度对比（相同输入，比较输出差异）
test_input = torch.randn(8, 4)
with torch.no_grad():
    torch_output = model(test_input).numpy()

ort_output = ort_session.run(None, {"input": to_numpy(test_input)})[0]

# 计算最大误差
max_diff = np.max(np.abs(torch_output - ort_output))
print(f"PyTorch与ONNX Runtime输出最大误差: {max_diff:.6e}")
assert max_diff < 1e-5, "精度差异过大，转换可能存在问题"
print("精度对比通过：两者输出一致")

# 5. 速度对比（各推理100次取平均）
inp_np = to_numpy(test_input)
inp_torch = test_input

# PyTorch推理耗时
with torch.no_grad():
    # 预热
    for _ in range(10):
        _ = model(inp_torch)
    start = time.perf_counter()
    for _ in range(100):
        _ = model(inp_torch)
    torch_time = (time.perf_counter() - start) / 100

# ONNX Runtime推理耗时
for _ in range(10):
    _ = ort_session.run(None, {"input": inp_np})
start = time.perf_counter()
for _ in range(100):
    _ = ort_session.run(None, {"input": inp_np})
ort_time = (time.perf_counter() - start) / 100

print(f"PyTorch平均推理耗时: {torch_time*1000:.3f} ms")
print(f"ONNX Runtime平均推理耗时: {ort_time*1000:.3f} ms")
print(f"加速比: {torch_time/ort_time:.2f}x")

运行方式

python onnx_inference_demo.py

验证方法与判定标准

精度验证：程序会计算PyTorch与ONNX Runtime输出的最大误差，若小于1e-5则打印“精度对比通过”，否则断言失败。这是转换正确性的核心指标。
速度验证：程序分别对PyTorch和ONNX Runtime执行100次推理并计算平均耗时，输出两者的毫秒耗时及加速比。加速比大于1说明ONNX Runtime在当前环境下更快。
结构验证：onnx.checker.check_model通过且打印“ONNX模型导出并检查通过”，说明导出的ONNX文件结构合法。

常见问题

加速比不明显或更慢：对于极小模型（如本例），ONNX Runtime的框架开销可能抵消优化收益。建议在真实业务模型（如ResNet、BERT等）上测试，加速效果更显著。可通过providers参数尝试CUDAExecutionProvider（需安装GPU版ONNX Runtime）。
动态维度报错：导出时已设置dynamic_axes，若遇到维度问题，检查输入是否满足最小维度要求。
算子不支持：若模型包含自定义算子，需在导出前注册或替换为ONNX支持的等价实现。