4、推理引擎入门：ONNX Runtime安装与使用、CPU推理与GPU推理、性能基准测试

各位同学，欢迎来到第四章。

前几章我们聊了模型训练、导出，算是把模型「生」出来了。但模型生出来不是供着的，得让它跑起来，对吧？

这一章，我们就来聊聊怎么让模型真正「动」起来。说白了，就是选一个趁手的推理引擎。

4.1 为什么是ONNX Runtime？

市面上推理引擎不少，TensorRT、OpenVINO、TNN……我为什么先讲ONNX Runtime？

原因很简单：通用性。

ONNX（Open Neural Network Exchange）本身是一个中间格式，就像一个「通用语言」。你不管用什么框架训练（PyTorch、TensorFlow），都能转成ONNX。而ONNX Runtime就是专门跑这种格式的引擎。

我个人习惯是：先转ONNX，再用ORT跑。这样后续换硬件、换平台，成本最低。

小提示： 我在项目中遇到过客户非要拿PyTorch直接部署的情况。结果环境依赖一堆，版本冲突搞了三天。后来我直接转ONNX，半小时搞定。你想想看，哪个更香？

4.2 安装ONNX Runtime

安装其实很简单。但要注意区分CPU版本和GPU版本。

4.2.1 CPU版本安装

如果你只是本地测试，或者服务器没有独立显卡，用CPU版本就够了。

pip install onnxruntime

嗯，就这一行。装完你可以验证一下：

import onnxruntime
print(onnxruntime.__version__)
print(onnxruntime.get_device())  # 输出应该是 'CPU'

4.2.2 GPU版本安装

要上GPU推理，得装带CUDA支持的版本。这里有个坑，我当年踩过。

pip install onnxruntime-gpu

注意：不是在同一个环境里装两个版本。CPU和GPU版本是互斥的。

警告： 我曾经在同一个conda环境里先后装了onnxruntime和onnxruntime-gpu，结果推理时一直调用CPU，GPU死活不工作。排查了半天才发现是版本冲突。所以，请新建一个干净环境。

4.3 CPU推理 vs GPU推理

装好了，我们来跑个实际例子。假设你有一个训练好的ResNet-50模型，已经转成了ONNX格式。

4.3.1 CPU推理代码

import onnxruntime as ort
import numpy as np

# 创建推理会话
session = ort.InferenceSession("resnet50.onnx")

# 构造输入数据（假设输入是224x224的图片）
input_name = session.get_inputs()[0].name
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 执行推理
outputs = session.run(None, {input_name: input_data})
print("推理完成，输出形状:", outputs[0].shape)

这段代码CPU和GPU都能跑。区别在哪？在于创建session时用的provider。

4.3.2 GPU推理代码

import onnxruntime as ort

# 显式指定使用CUDA
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
session = ort.InferenceSession("resnet50.onnx", providers=providers)

# 后面的代码和CPU版本一样
input_name = session.get_inputs()[0].name
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = session.run(None, {input_name: input_data})

看到没？核心区别就是providers参数。我建议你总是把CPUExecutionProvider放在最后，作为兜底。

核心要点： GPU推理不是「自动」的。你必须显式指定CUDA provider，否则ORT默认用CPU。

4.4 性能基准测试

光说GPU比CPU快，那快多少？咱们得用数据说话。

4.4.1 写一个简单的测速脚本

import time
import onnxruntime as ort
import numpy as np

def benchmark(model_path, providers, num_runs=100):
    session = ort.InferenceSession(model_path, providers=providers)
    input_name = session.get_inputs()[0].name
    input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
    
    # 预热（warm up）
    for _ in range(10):
        session.run(None, {input_name: input_data})
    
    # 正式测试
    start = time.time()
    for _ in range(num_runs):
        session.run(None, {input_name: input_data})
    end = time.time()
    
    avg_time = (end - start) / num_runs
    fps = 1.0 / avg_time
    return avg_time, fps

# CPU测试
cpu_time, cpu_fps = benchmark("resnet50.onnx", ['CPUExecutionProvider'])
print(f"CPU: 平均耗时 {cpu_time*1000:.2f} ms, FPS: {cpu_fps:.2f}")

# GPU测试
gpu_time, gpu_fps = benchmark("resnet50.onnx", ['CUDAExecutionProvider', 'CPUExecutionProvider'])
print(f"GPU: 平均耗时 {gpu_time*1000:.2f} ms, FPS: {gpu_fps:.2f}")

4.4.2 典型结果对比

我在一台普通工作站上跑过（CPU: i7-12700, GPU: RTX 3060），结果大概是这样：

推理设备 平均耗时（ms） FPS 加速比
CPU（12核） 45.2 22.1 1x
GPU（RTX 3060） 3.8 263.2 约12x

看到了吧？GPU推理快了整整一个数量级。这就是为什么生产环境里，能上GPU就上GPU。

我的经验： 如果你用的是小模型（比如MobileNet），CPU和GPU差距可能没那么大，甚至CPU更快（因为GPU启动有开销）。但一旦模型变大，GPU的优势就碾压了。所以，先测再决定。

4.5 避坑指南

最后，分享几个我实际踩过的坑：

  • CUDA版本不匹配： ONNX Runtime的GPU版本对CUDA版本有严格要求。比如ort 1.15要求CUDA 11.8。装错了直接报错。我建议你装之前先查官方文档的兼容性矩阵。
  • 显存泄漏： 如果你在循环里反复创建session，显存会一直涨。正确做法是只创建一次session，重复使用。
  • 动态输入形状： 有些模型输入尺寸不固定，ORT默认不支持。需要在导出ONNX时设置dynamic_axes参数。这个我们后面章节会细讲。

好了，这一章的内容就到这里。下一章我们会深入ONNX模型的导出细节，包括怎么处理动态输入、怎么优化算子。到时候见。

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