3、ONNX Runtime基础：ONNX Runtime架构、安装与配置、Python API快速上手

好，咱们进入ONNX Runtime的世界。说实话，ONNX Runtime是我在生产环境中用得最多的推理引擎之一。它不像TensorRT那样只认NVIDIA显卡，也不像OpenVINO那样绑定Intel硬件。ONNX Runtime的目标很纯粹——让你训练好的模型，能在各种设备上跑起来，而且跑得尽量快。

我个人习惯把ONNX Runtime看作一个“模型翻译官”。你给它一个ONNX模型，它负责调用底层各种优化器，最后在CPU、GPU甚至NPU上执行。你不需要关心底层是CUDA还是DirectML，它帮你搞定。

ONNX Runtime架构：它到底是怎么工作的？

先看整体架构。ONNX Runtime分三层：

  • 前端（Frontend）：负责解析ONNX模型，做图优化。比如常量折叠、算子融合。这些优化不依赖硬件。
  • 执行引擎（Execution Engine）：核心调度器。它决定算子按什么顺序执行，怎么分配内存。
  • 后端（Backend / Providers）：真正干活的地方。CPU用OpenMP或MKL，GPU用CUDA Execution Provider，Intel芯片用OpenVINO Execution Provider。

为什么会这样分层？说白了，就是为了解耦。你换硬件，只需要换后端Provider，前端优化逻辑完全不用动。我在项目中遇到过客户从NVIDIA换到AMD显卡，我们只需要切换Execution Provider，模型代码一行没改。

核心要点：ONNX Runtime的架构设计，本质上是“一次优化，到处运行”。前端做通用优化，后端做硬件特化。

安装与配置：别踩这些坑

安装其实很简单，但有几个细节我建议你注意。

CPU版本安装：

pip install onnxruntime

GPU版本安装：

pip install onnxruntime-gpu

嗯，这里要注意。GPU版本对CUDA和cuDNN版本有严格对应关系。我曾经因为CUDA 11.8配了onnxruntime-gpu 1.15，结果跑起来直接报错。后来查文档才发现，1.15要求CUDA 11.6。你想想看，这种版本错配问题排查起来多浪费时间。

我建议你安装前先看一眼官方版本兼容表：

ONNX Runtime版本CUDA版本cuDNN版本
1.15.x11.68.4
1.16.x11.88.6
1.17.x12.08.9

避坑指南：我曾经在服务器上同时装了多个CUDA版本，结果onnxruntime-gpu死活找不到正确的CUDA库。后来我改用Docker镜像，把CUDA版本锁死，再也没出过问题。如果你也遇到类似问题，试试Docker。

验证安装是否成功：

import onnxruntime
print(onnxruntime.get_device())  # 输出 CPU 或 GPU
print(onnxruntime.get_available_providers())  # 列出可用后端

如果输出里没有CUDA Execution Provider，说明GPU版本没装对。别慌，检查一下CUDA环境变量。

Python API快速上手：三行代码跑推理

好，咱们直接上手写代码。ONNX Runtime的Python API非常简洁，核心就三步：

  1. 创建推理会话（InferenceSession）
  2. 获取输入输出名称
  3. 运行推理

看代码：

import onnxruntime
import numpy as np

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

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

# 3. 构造输入数据（假设模型输入是1x3x224x224）
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

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

print("推理完成，输出形状：", result[0].shape)

是不是很简单？但这里有个细节——输入数据的dtype必须是float32。我见过有人用float64传进去，结果推理速度慢了好几倍。ONNX Runtime默认优化是针对float32的，你用float64等于绕过了所有优化。

个人经验：我习惯在推理前加一句 input_data = input_data.astype(np.float32)，确保类型正确。这个习惯帮我避免了很多线上bug。

配置会话选项：让推理更灵活

实际项目中，你不可能只用默认配置。比如你想指定用GPU跑，或者限制线程数。这时候需要配置SessionOptions。

import onnxruntime

# 创建配置对象
options = onnxruntime.SessionOptions()

# 设置线程数
options.intra_op_num_threads = 4
options.inter_op_num_threads = 2

# 设置日志级别（0:Verbose, 1:Info, 2:Warning, 3:Error）
options.log_severity_level = 3

# 启用/禁用内存优化
options.enable_mem_pattern = True

# 指定Execution Provider
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']

# 创建会话
session = onnxruntime.InferenceSession(
    "model.onnx",
    sess_options=options,
    providers=providers
)

这里我解释一下两个线程参数的区别：

  • intra_op_num_threads：单个算子内部的并行线程数。比如矩阵乘法，用多线程加速。
  • inter_op_num_threads：不同算子之间的并行线程数。如果两个算子没有依赖关系，可以并行执行。

我个人习惯把intra_op设置成CPU核心数，inter_op设置成2。你想想看，如果inter_op设太大，线程切换开销反而会拖慢速度。

性能调优：几个实用技巧

最后分享几个我在项目中用过的调优技巧：

  1. 启用Graph Optimization：ONNX Runtime默认开启所有优化，但你可以手动控制级别。
options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
  1. 使用IO Binding减少内存拷贝：如果你在GPU上推理，用IO Binding可以避免CPU和GPU之间的数据搬运。
# 创建IO Binding
io_binding = session.io_binding()

# 绑定GPU输入输出
io_binding.bind_cpu_input(input_name, input_data)
io_binding.bind_output(output_name, device_type='cuda')

# 运行
session.run_with_iobinding(io_binding)

# 获取结果
result = io_binding.copy_outputs_to_cpu()[0]
  1. 预热（Warm-up）：第一次推理通常比较慢，因为要加载模型、分配内存。我建议在正式推理前，先跑一次假数据做预热。
# 预热
_ = session.run([output_name], {input_name: dummy_input})
print("预热完成")

总结一下：ONNX Runtime的核心价值在于“一次导出，到处运行”。你只需要掌握三个API——InferenceSession、run、SessionOptions，就能应对90%的部署场景。剩下的10%，靠经验和文档慢慢摸索。

好，这一章就到这里。下一章咱们聊聊ONNX Runtime的高级用法，包括自定义算子、量化推理和动态形状处理。这些东西在实际项目中非常实用，尤其是当你遇到模型不兼容或者性能瓶颈的时候。

← 上一章 📖 返回目录 下一章 →