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原理

说白了就三步:

  1. 按置信度排序,选最高的框
  2. 计算这个框和其他框的IoU
  3. 去掉IoU大于阈值的框(认为它们检测的是同一个目标)
  4. 重复,直到没有框剩下

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部署的完整流程串起来了:

ONNX Runtime部署流程 ONNX Runtime部署完整流程 1. ONNX Runtime安装 2. 模型加载与Session创建 3. 执行推理 4. 输出张量解析 5. NMS后处理 关键配置 • 线程数: 核心数-1 • 优化级别: ORT_ENABLE_ALL • 提供者: CPUExecutionProvider 常见坑点 • 输入形状不匹配 • 坐标归一化范围 • NMS阈值调优 输出: 最终检测框列表 [x1,y1,x2,y2,conf,class_id]

4.6 性能对比

不同后处理方式的性能差异,我实测过:

后处理方式 单帧耗时 (ms) 适用场景
纯Python NMS 15-25 原型验证、框数量少
NumPy向量化NMS 5-10 嵌入式实时推理
C++扩展NMS 1-3 高性能要求
硬件加速NMS (NPU) <1 专用芯片

注意: 表格中的数据是在树莓派4B上测的。不同设备差异很大。我建议你在目标设备上重新测一遍,别直接套用别人的数据。

好了,ONNX Runtime部署的核心流程就这些。从安装到NMS,每一步都有细节。你动手跑一遍,遇到问题再回来看,理解会更深刻。