1. YOLOv8n模型ONNX导出与验证

任务目标

将YOLOv8n目标检测模型导出为ONNX格式，并验证导出模型的输出张量形状与数值正确性，确保ONNX模型与PyTorch原始模型在相同输入下推理结果一致。

环境依赖

需要安装以下依赖：

pip install ultralytics onnx onnxruntime

完整实现代码

import torch
import numpy as np
import onnx
import onnxruntime as ort
from ultralytics import YOLO

# 1. 加载YOLOv8n预训练模型
model = YOLO('yolov8n.pt')

# 2. 导出为ONNX格式（动态batch，固定输入尺寸640x640）
export_path = model.export(format='onnx', imgsz=640, dynamic=True, simplify=True)
print(f'ONNX模型已保存至: {export_path}')

# 3. 构造测试输入（与训练预处理一致：RGB、0-255、BCHW）
dummy_input = np.random.randint(0, 255, size=(1, 3, 640, 640), dtype=np.uint8).astype(np.float32)

# 4. 获取PyTorch模型输出（作为基准）
model.eval()
with torch.no_grad():
    torch_output = model.predict(source=dummy_input, imgsz=640, verbose=False)
    # 取第一个检测结果的原始输出张量（形状: [1, 84, 8400]）
    torch_preds = torch_output[0].boxes.data.cpu().numpy()  # [N, 6] 格式: x1,y1,x2,y2,conf,cls

# 5. 使用ONNX Runtime进行推理
ort_session = ort.InferenceSession(export_path)
onnx_input = {ort_session.get_inputs()[0].name: dummy_input}
onnx_output = ort_session.run(None, onnx_input)[0]  # 形状: [1, 84, 8400]

# 6. 验证输出张量形状
print(f'ONNX输出形状: {onnx_output.shape}')
assert onnx_output.shape == (1, 84, 8400), f'形状不匹配: {onnx_output.shape}'

# 7. 验证数值正确性（将ONNX输出转换为与PyTorch相同的检测框格式）
# ONNX输出格式: [batch, 4+80, num_anchors] = [1, 84, 8400]
# 需要解码为xyxy格式并与PyTorch结果对比
def decode_onnx_output(output, conf_thres=0.25):
    """将ONNX原始输出解码为检测框列表"""
    output = output[0]  # [84, 8400]
    num_classes = 80
    boxes = []
    for i in range(output.shape[1]):
        scores = output[4:, i]
        class_id = np.argmax(scores)
        confidence = scores[class_id]
        if confidence < conf_thres:
            continue
        cx, cy, w, h = output[:4, i]
        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 np.array(boxes) if boxes else np.empty((0, 6))

onnx_boxes = decode_onnx_output(onnx_output)

# 8. 对比检测结果（允许坐标微小误差）
if len(onnx_boxes) > 0 and len(torch_preds) > 0:
    # 按置信度排序后取前5个框对比
    onnx_boxes = onnx_boxes[onnx_boxes[:, 4].argsort()[::-1]][:5]
    torch_preds = torch_preds[torch_preds[:, 4].argsort()[::-1]][:5]
    
    print(f'PyTorch检测框数量: {len(torch_preds)}, ONNX检测框数量: {len(onnx_boxes)}')
    print('前5个检测框对比（x1,y1,x2,y2,conf,cls）:')
    for i in range(min(5, len(onnx_boxes), len(torch_preds))):
        diff = np.abs(onnx_boxes[i] - torch_preds[i])
        print(f'框{i+1} 坐标最大差异: {diff[:4].max():.4f}, 置信度差异: {diff[4]:.6f}')
        assert diff[:4].max() < 1.0, f'坐标差异过大: {diff}'
        assert diff[4] < 0.01, f'置信度差异过大: {diff}'
    print('数值验证通过：ONNX与PyTorch输出一致')
else:
    print('警告：未检测到足够的目标框，无法进行数值对比')
    print('请更换测试输入或降低置信度阈值后重试')

运行与验证方法

运行步骤：

  1. 确保已安装依赖并下载yolov8n.pt预训练权重（首次运行会自动下载）
  2. 直接执行上述Python脚本
  3. 观察控制台输出

通过条件：

  • 脚本无异常退出，打印“ONNX模型已保存至”路径
  • 输出张量形状断言通过，打印ONNX输出形状: (1, 84, 8400)
  • 若检测到目标框，坐标最大差异小于1像素且置信度差异小于0.01，打印“数值验证通过”
  • 若未检测到目标框，脚本会给出警告提示，此时应更换输入图像或降低置信度阈值重新验证

关键说明

ONNX导出的输出为原始预测张量[1, 84, 8400]，其中84=4（边界框坐标）+80（COCO类别数），8400为不同尺度特征图上的锚点总数。PyTorch的model.predict()返回的是经过NMS后处理的检测结果，因此验证时需要将ONNX输出解码为相同格式再对比。本案例使用随机噪声作为输入，仅用于验证模型结构正确性和数值一致性，实际部署时应使用真实图像测试。

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