4、ONNX Runtime部署：安装配置、推理与性能调优

好，咱们今天聊聊ONNX Runtime。说实话，这玩意儿在我日常工作中，几乎是绕不开的一个工具。你想想看，不管你是用PyTorch、TensorFlow还是其他框架训练出来的模型，最后想落地到生产环境，总得有个统一的“中间人”吧？ONNX Runtime就是干这个的。

我个人习惯，把ONNX Runtime看作是一个“模型翻译官”。它把各种框架的模型，都转成ONNX这种通用格式，然后用自己的高性能引擎去跑。这样，你就不用为每个框架单独写一套推理代码了，省心不少。

4.1 ONNX Runtime的安装与配置

安装这块，其实挺简单的。但我在项目中遇到过一个问题：很多人直接 pip install onnxruntime，结果发现跑起来特别慢。为什么？因为默认装的是CPU版本。

如果你有NVIDIA显卡，一定要装GPU版本。命令是这样的：

# CPU版本
pip install onnxruntime

# GPU版本（需要CUDA和cuDNN）
pip install onnxruntime-gpu

安装完成后，你可以用一行代码验证一下：

import onnxruntime
print(onnxruntime.get_device())

如果输出是 GPU，那就对了。如果是 CPU，嗯，你可能装错了版本。

4.2 使用ONNX Runtime进行推理

推理流程其实就三步：加载模型、准备输入、运行推理。说白了，跟用其他框架差不多。

我拿一个图像分类模型举个例子。假设你已经有了一个 model.onnx 文件：

import onnxruntime
import numpy as np
from PIL import Image

# 1. 创建推理会话
session = onnxruntime.InferenceSession("model.onnx")

# 2. 获取输入输出信息
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 3. 准备输入数据（假设是224x224的图片）
img = Image.open("cat.jpg").resize((224, 224))
img_array = np.array(img).astype(np.float32) / 255.0
img_array = np.transpose(img_array, (2, 0, 1))  # HWC -> CHW
img_array = np.expand_dims(img_array, axis=0)   # 增加batch维度

# 4. 运行推理
results = session.run([output_name], {input_name: img_array})

# 5. 处理结果
predicted_class = np.argmax(results[0])
print(f"预测类别: {predicted_class}")

你看，代码量其实很少。但这里有个坑：输入数据的格式必须跟模型训练时完全一致。包括数据类型、形状、归一化方式。我见过有人把 float32 传成了 int8，结果推理结果全乱了。

4.3 性能调优与算子兼容性

这部分是重点，也是我踩坑最多的地方。性能调优，说白了就是让模型跑得更快、占内存更少。

4.3.1 性能调优三板斧

第一板斧：开启优化选项。 ONNX Runtime默认就带了一些优化，但你可以手动开启更多：

import onnxruntime

# 开启所有优化
options = onnxruntime.SessionOptions()
options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL

# 设置执行模式（并行执行）
options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL

session = onnxruntime.InferenceSession("model.onnx", options)

第二板斧：使用TensorRT加速。 如果你用的是NVIDIA显卡，强烈建议开启TensorRT。我试过，推理速度能提升2-5倍。

import onnxruntime

options = onnxruntime.SessionOptions()
options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL

# 配置TensorRT
trt_options = onnxruntime.SessionOptions()
trt_options.enable_trt = True
trt_options.trt_max_workspace_size = 1 << 30  # 1GB

session = onnxruntime.InferenceSession("model.onnx", trt_options)

第三板斧：调整线程数。 这个在CPU上特别管用。你可以根据CPU核心数来设置：

options = onnxruntime.SessionOptions()
options.intra_op_num_threads = 4  # 内部操作线程数
options.inter_op_num_threads = 2  # 并行操作线程数

我个人习惯，先设成CPU核心数的一半，然后慢慢往上调，直到性能不再提升为止。

4.3.2 算子兼容性

算子兼容性，说白了就是你的模型里用的那些操作，ONNX Runtime能不能支持。我遇到过最头疼的问题，就是模型导出时一切正常，但一用ONNX Runtime推理就报错。

为什么会这样？因为有些算子，比如 torch.nn.functional.grid_sample，在ONNX里可能没有对应的实现。或者实现方式不一样。

我建议你这样做：

1. 导出前检查算子： 用 torch.onnx.export 时，设置 verbose=True，可以看到每个算子被映射成了什么。
2. 使用ONNX Runtime的兼容性检查工具：

import onnx
from onnxruntime.transformers.onnx_model import OnnxModel

model = onnx.load("model.onnx")
onnx_model = OnnxModel(model)

# 检查不支持的算子
unsupported_ops = onnx_model.get_unsupported_ops()
print(f"不支持的算子: {unsupported_ops}")

如果发现有不支持的算子，你有几个选择：

• 替换算子： 在模型里用更通用的操作替代。比如用 resize 替代 grid_sample。
• 自定义算子： 如果实在绕不开，可以自己写一个自定义算子。但说实话，这个比较麻烦，我一般不建议。
• 升级ONNX Runtime版本： 新版本通常会支持更多算子。我遇到过 LayerNormalization 在旧版本不支持，升级后就解决了。

4.3.3 性能对比小实验

最后，我分享一个我常用的性能对比方法。你可以用下面的代码，对比不同配置下的推理速度：

import time
import onnxruntime
import numpy as np

def benchmark(session, input_data, iterations=100):
    # 预热
    for _ in range(10):
        session.run(None, {input_name: input_data})
    
    # 正式测试
    start = time.time()
    for _ in range(iterations):
        session.run(None, {input_name: input_data})
    end = time.time()
    
    avg_time = (end - start) / iterations * 1000  # 毫秒
    print(f"平均推理时间: {avg_time:.2f} ms")

# 测试不同配置
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 默认配置
session_default = onnxruntime.InferenceSession("model.onnx")
print("默认配置:")
benchmark(session_default, input_data)

# 优化配置
options = onnxruntime.SessionOptions()
options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
session_optimized = onnxruntime.InferenceSession("model.onnx", options)
print("优化配置:")
benchmark(session_optimized, input_data)

嗯，这个对比结果，往往能让你直观地看到优化的效果。我见过最夸张的一次，优化后推理时间从200ms降到了30ms，整整快了6倍多。

好了，关于ONNX Runtime的部署，我就讲这么多。记住，安装时注意版本，推理时注意数据格式，调优时多试试不同配置。踩坑不可怕，关键是踩完坑要记得填上。