3、ONNX Runtime部署：安装配置与推理实战

说实话，ONNX Runtime 是我个人最常用的推理引擎之一。为什么？因为它省心。你训练好的 PyTorch 或 TensorFlow 模型，转成 ONNX 格式后，ONNX Runtime 基本都能接住。我最早接触它是在一个工业质检项目里，客户要求模型必须跑在 Windows 的 C++ 环境里，还不能装 Python。当时我就靠 ONNX Runtime 搞定的。

这一章，咱们就聊聊 ONNX Runtime 的安装、配置，以及 Python 和 C++ 两种推理接口怎么用。还有动态轴、静态轴的处理，以及算子兼容性排查——这些坑我基本都踩过一遍。

3.1 ONNX Runtime 安装与配置

安装其实很简单。但不同平台、不同硬件，细节上还是有区别的。

3.1.1 Python 安装

如果你只是做快速验证，用 pip 装 CPU 版本就够了：

pip install onnxruntime

要是想用 GPU 加速，记得装 GPU 版本：

pip install onnxruntime-gpu

3.1.2 C++ 安装

C++ 部署稍微麻烦一点。你可以从 GitHub Release 页面下载预编译的库，也可以自己编译。

我个人习惯用 vcpkg 来管理：

vcpkg install onnxruntime

或者用 CMake 直接 FetchContent：

FetchContent_Declare(
  onnxruntime
  GIT_REPOSITORY https://github.com/microsoft/onnxruntime.git
  GIT_TAG v1.15.0
)
FetchContent_MakeAvailable(onnxruntime)

3.2 Python 推理接口

Python 接口非常直观。说白了就是三步：创建会话、准备输入、运行推理。

import onnxruntime as ort
import numpy as np

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

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

# 3. 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 4. 运行推理
output = sess.run([output_name], {input_name: input_data})
print(output[0].shape)

嗯，这里要注意一点：ONNX Runtime 默认输入是 float32 类型。如果你传了 float64，它会报错。我刚开始用的时候就被这个坑过。

3.2.1 设置运行选项

你可以通过 SessionOptions 来配置运行行为：

options = ort.SessionOptions()
options.intra_op_num_threads = 4  # 设置线程数
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

sess = ort.InferenceSession("model.onnx", options)

3.3 C++ 推理接口

C++ 接口比 Python 稍微复杂一点，但性能更好。尤其是在生产环境中，C++ 是首选。

#include <onnxruntime/core/session/onnxruntime_cxx_api.h>
#include <vector>

int main() {
    // 1. 创建环境
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "test");

    // 2. 创建会话选项
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(4);

    // 3. 加载模型
    Ort::Session session(env, "model.onnx", session_options);

    // 4. 准备输入
    std::vector<int64_t> input_shape = {1, 3, 224, 224};
    std::vector<float> input_data(1*3*224*224, 1.0f);

    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(
        OrtArenaAllocator, OrtMemTypeDefault);

    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        memory_info, input_data.data(), input_data.size(),
        input_shape.data(), input_shape.size());

    // 5. 运行推理
    const char* input_names[] = {"input"};
    const char* output_names[] = {"output"};

    auto output_tensors = session.Run(Ort::RunOptions{nullptr},
                                      input_names, &input_tensor, 1,
                                      output_names, 1);

    // 6. 获取输出
    float* output_data = output_tensors[0].GetTensorMutableData<float>();
    auto output_shape = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape();

    return 0;
}

3.4 动态轴与静态轴处理

这是 ONNX 部署中最容易出问题的地方。什么叫动态轴？说白了就是模型输入的形状可以变化。比如 NLP 模型，句子长度不固定，这就是动态轴。

3.4.1 静态轴

静态轴就是输入形状完全固定。比如图像分类模型，输入永远是 (1, 3, 224, 224)。这种情况下，ONNX Runtime 会做很多优化，性能最好。

3.4.2 动态轴

动态轴需要你在导出 ONNX 时指定。以 PyTorch 为例：

import torch

class MyModel(torch.nn.Module):
    def forward(self, x):
        return x * 2

model = MyModel()
dummy_input = torch.randn(1, 3, 224, 224)

# 动态轴：batch size 和 height, width 都可以变
dynamic_axes = {
    'input': {0: 'batch_size', 2: 'height', 3: 'width'},
    'output': {0: 'batch_size', 2: 'height', 3: 'width'}
}

torch.onnx.export(
    model,
    dummy_input,
    "dynamic_model.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes=dynamic_axes
)

3.4.3 在 ONNX Runtime 中处理动态轴

Python 端处理动态轴很简单，直接传不同形状的输入就行：

# 动态轴模型，可以传不同 batch size
input1 = np.random.randn(1, 3, 224, 224).astype(np.float32)
output1 = sess.run([output_name], {input_name: input1})

input2 = np.random.randn(4, 3, 224, 224).astype(np.float32)
output2 = sess.run([output_name], {input_name: input2})

C++ 端也一样，只要在创建 Tensor 时指定不同的 shape 即可。

3.5 算子兼容性排查

这是最让人头疼的部分。你模型在 PyTorch 里跑得好好的，转成 ONNX 后，某些算子不支持，直接报错。

3.5.1 常见的不兼容算子

算子	问题描述	解决方案
torch.where	ONNX opset 较低时不支持	升级 opset 或用 torch.where 替代
torch.nn.functional.grid_sample	ONNX 标准算子集中没有	使用自定义算子或替换实现
torch.einsum	部分 einsum 模式不支持	拆解成多个 matmul 和 transpose
torch.topk	某些 opset 中 k 必须是常量	导出时固定 k 值

3.5.2 排查方法

我一般用三步法来排查算子兼容性：

1. 先用 onnxruntime 直接跑，看报错信息。ONNX Runtime 的报错通常很明确，会告诉你哪个算子不支持。
2. 用 onnx.checker 检查模型：

import onnx
model = onnx.load("model.onnx")
onnx.checker.check_model(model)
print("模型检查通过")

3. 用 onnxruntime 的推理能力检查：

import onnxruntime as ort

# 查看支持的算子
providers = ort.get_available_providers()
print("可用 provider:", providers)

# 查看某个 opset 支持的算子
opsets = ort.get_all_providers()
for p in opsets:
    print(p)

3.5.3 自定义算子

如果标准算子集里确实没有，那就只能自己写自定义算子。不过这个比较复杂，我建议你先看看有没有替代方案。比如 grid_sample，可以用多个卷积和插值操作来模拟。

实在不行，再考虑自定义算子。ONNX Runtime 支持通过 C++ 扩展来注册自定义算子。具体做法是：

1. 实现算子的前向计算逻辑
2. 注册到 ONNX Runtime 的算子库中
3. 在导出 ONNX 时使用这个自定义算子

嗯，这个内容比较深，咱们课程后面会有专门章节来讲。这里先提个醒。

3.6 本章知识体系

下面这张图是我画的，把本章的核心逻辑串起来了：

这张图把本章的核心内容都串起来了。从安装配置开始，到 Python 和 C++ 两种推理接口，再到动态轴静态轴的处理，最后是算子兼容性排查。你跟着这个脉络走，基本不会迷路。

---

好了，ONNX Runtime 部署这部分就聊到这儿。说实话，ONNX Runtime 是我用过的最省心的推理引擎之一。只要模型导出没问题，部署基本就是复制粘贴代码的事。但要注意，动态轴和算子兼容性这两个坑，一定要提前排查。我当年就是没注意动态轴的问题，上线后才发现推理速度慢得离谱，后来改回静态轴才解决。

下一章咱们聊聊 TensorRT 部署，那个更刺激，性能也更好。但坑也更多。到时候我把我踩过的坑都告诉你。