4. ONNX Runtime部署:从安装到后处理全流程
ONNX Runtime,说白了就是微软开源的推理引擎。它能把你训练好的模型,在各种硬件上跑起来。我最早接触它是在一个工业检测项目里,客户要求模型必须能在Windows和Linux上无缝切换——ONNX Runtime正好解决了这个痛点。
这一章,我们一步步把ONNX Runtime部署的完整流程走通。从安装开始,到模型加载、推理,再到输出解析和后处理。嗯,每一步我都会把项目中踩过的坑告诉你。
4.1 ONNX Runtime安装
安装其实很简单。但不同平台、不同硬件,选对版本很重要。
| 平台 | 安装命令 | 说明 |
|---|---|---|
| CPU (通用) | pip install onnxruntime |
最轻量,无硬件加速 |
| GPU (CUDA) | pip install onnxruntime-gpu |
需要CUDA和cuDNN |
| ARM (树莓派等) | pip install onnxruntime |
官方支持ARM架构 |
| NPU/TPU | 厂商定制版 | 如Rockchip、华为昇腾 |
我的建议: 嵌入式端优先用CPU版本。GPU版本依赖太多,交叉编译时容易出问题。我之前在RK3588上折腾GPU版,折腾了两天才发现官方推荐用CPU版+NPU协同推理。
验证安装是否成功:
import onnxruntime
print(onnxruntime.__version__)
print(onnxruntime.get_device()) # 输出 CPU 或 GPU
如果打印出版本号,说明安装成功。注意,get_device()返回的是当前可用的设备,不是你想用的设备。
4.2 模型加载与推理
模型加载,核心就一个API:onnxruntime.InferenceSession。但参数配置有讲究。
4.2.1 创建Session
import onnxruntime as ort
# 基本加载
session = ort.InferenceSession("yolov8n.onnx")
# 带配置加载(推荐)
options = ort.SessionOptions()
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
options.intra_op_num_threads = 4 # 根据核心数调整
session = ort.InferenceSession(
"yolov8n.onnx",
sess_options=options,
providers=['CPUExecutionProvider']
)
我曾经踩过的坑: 不设置线程数,默认会用满所有核心。在嵌入式设备上,这会导致其他任务卡死。我建议根据设备核心数,留一个核心给系统用。比如4核设备,设3个线程就够了。
4.2.2 获取输入输出信息
推理前,先搞清楚模型的输入输出长什么样:
# 获取输入信息
input_name = session.get_inputs()[0].name
input_shape = session.get_inputs()[0].shape
input_type = session.get_inputs()[0].type
# 获取输出信息
output_name = session.get_outputs()[0].name
output_shape = session.get_outputs()[0].shape
print(f"输入: {input_name}, 形状: {input_shape}, 类型: {input_type}")
print(f"输出: {output_name}, 形状: {output_shape}")
YOLOv8的输入通常是 [1, 3, 640, 640],输出是 [1, 84, 8400]。为什么是84?80个类别 + 4个坐标。8400是检测框的数量。
4.2.3 执行推理
import numpy as np
# 假设 input_tensor 已经是预处理好的 [1,3,640,640] 的float32数组
results = session.run(
[output_name], # 输出名称列表
{input_name: input_tensor} # 输入字典
)
# results[0] 的形状: (1, 84, 8400)
output = results[0]
推理就这么简单。但真正的难点在后面——怎么把这一堆数字变成我们能理解的检测框。
4.3 输入输出张量解析
YOLOv8的输出格式,和YOLOv5不太一样。我第一次迁移时就被坑了。
4.3.1 输出结构拆解
输出张量 [1, 84, 8400] 的含义:
- 维度0 (1): batch size,通常为1
- 维度1 (84): 每个检测框的特征向量。前4个是框坐标 (cx, cy, w, h),后80个是类别置信度
- 维度2 (8400): 检测框数量。来自三个尺度的特征图融合
解析代码:
def parse_output(output, conf_threshold=0.5):
"""
output: numpy array, shape (1, 84, 8400)
返回: 检测框列表 [ [x1,y1,x2,y2,conf,class_id], ... ]
"""
# 转置为 (8400, 84)
output = output.transpose((0, 2, 1)) # (1, 8400, 84)
output = output[0] # (8400, 84)
boxes = []
for i in range(output.shape[0]):
# 前4个是框坐标
cx, cy, w, h = output[i, :4]
# 第5到84是类别置信度
class_scores = output[i, 4:]
class_id = np.argmax(class_scores)
confidence = class_scores[class_id]
if confidence < conf_threshold:
continue
# 转换为中心坐标 -> 左上角右下角
x1 = cx - w / 2
y1 = cy - h / 2
x2 = cx + w / 2
y2 = cy + h / 2
boxes.append([x1, y1, x2, y2, confidence, class_id])
return boxes
关键点: 坐标是归一化的 [0,1] 范围。要映射到原图尺寸,需要乘以图像的宽和高。我习惯在解析时就做映射,省得后面再处理。
4.4 后处理:NMS实现
NMS(非极大值抑制)是目标检测的最后一道关卡。它的作用:去掉重复的检测框,只保留最好的那个。
4.4.1 NMS原理
说白了就三步:
- 按置信度排序,选最高的框
- 计算这个框和其他框的IoU
- 去掉IoU大于阈值的框(认为它们检测的是同一个目标)
- 重复,直到没有框剩下
4.4.2 手写NMS代码
def iou(box1, box2):
"""计算两个框的IoU"""
x1 = max(box1[0], box2[0])
y1 = max(box1[1], box2[1])
x2 = min(box1[2], box2[2])
y2 = min(box1[3], box2[3])
inter_area = max(0, x2 - x1) * max(0, y2 - y1)
box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
union_area = box1_area + box2_area - inter_area
return inter_area / union_area if union_area > 0 else 0
def nms(boxes, iou_threshold=0.5):
"""
boxes: [[x1,y1,x2,y2,conf,class_id], ...]
返回: 筛选后的框列表
"""
if len(boxes) == 0:
return []
# 按置信度降序排序
boxes = sorted(boxes, key=lambda x: x[4], reverse=True)
keep = []
while boxes:
# 取置信度最高的框
best = boxes.pop(0)
keep.append(best)
# 过滤掉与best重叠过多的框
boxes = [box for box in boxes
if iou(best[:4], box[:4]) < iou_threshold]
return keep
性能优化建议: 嵌入式设备上,纯Python的NMS可能成为瓶颈。我建议用NumPy向量化实现,或者用C++扩展。如果框数量不多(<1000),Python版本也够用。
4.4.3 完整推理流程
把前面所有步骤串起来:
def detect(image, session, conf_threshold=0.5, iou_threshold=0.5):
# 1. 预处理
input_tensor = preprocess(image) # 假设已实现
# 2. 推理
output = session.run(None, {session.get_inputs()[0].name: input_tensor})[0]
# 3. 解析输出
boxes = parse_output(output, conf_threshold)
# 4. NMS
final_boxes = nms(boxes, iou_threshold)
# 5. 映射到原图坐标
h, w = image.shape[:2]
for box in final_boxes:
box[0] = int(box[0] * w)
box[1] = int(box[1] * h)
box[2] = int(box[2] * w)
box[3] = int(box[3] * h)
return final_boxes
4.5 本章知识结构
下面这张图,把ONNX Runtime部署的完整流程串起来了:
4.6 性能对比
不同后处理方式的性能差异,我实测过:
| 后处理方式 | 单帧耗时 (ms) | 适用场景 |
|---|---|---|
| 纯Python NMS | 15-25 | 原型验证、框数量少 |
| NumPy向量化NMS | 5-10 | 嵌入式实时推理 |
| C++扩展NMS | 1-3 | 高性能要求 |
| 硬件加速NMS (NPU) | <1 | 专用芯片 |
注意: 表格中的数据是在树莓派4B上测的。不同设备差异很大。我建议你在目标设备上重新测一遍,别直接套用别人的数据。
好了,ONNX Runtime部署的核心流程就这些。从安装到NMS,每一步都有细节。你动手跑一遍,遇到问题再回来看,理解会更深刻。