Phase A · Step 3:模型导出 ONNX

目标

将 PC 端的模型(Pipeline 验证阶段用预训练权重 yolo26n.pt,训练完成后用 best.pt) 导出为固定 shape 的 ONNX 模型,并通过 ONNX checker 和 simplifier 验证, 确保能被板端 ATC 工具接受。

阶段定位

本步骤在 Phase A 的位置:

bash 复制代码
[01 PC 环境准备]
      ↓
[02 预训练权重与 Pipeline 验证准备]  ← 产生 weights/yolo26n.pt
      ↓
[03 模型导出 ONNX]                  ← 本步骤
      ↓
[04 ???]
      ↓
[05 ATC 模型转换] → [06 ACL 推理部署] → [07 端到端验证]

执行时机有两种:

  1. Pipeline 验证阶段 (推荐先做 ,1-2 小时搞定):用 02 准备好的 weights/yolo26n.pt 导出 ONNX验证 ATC / ACL 整条链路是否跑通即使 YOLO26 官方权重检测不出无人机也没关系,能跑出非空结果即视为 pipeline 通过
  2. 真实训练阶段 (数据集准备好后,09 训练完成再回来重跑本步骤):把 model= 参数改成 runs/detect/yolo26_drone_n640/weights/best.pt其它参数完全不变

两次执行的命令几乎一样,只有 model= 不同。

前提条件

  • 已完成 02 预训练权重与 Pipeline 验证准备
  • yolo26 conda 环境已激活
  • 已安装 ultralytics、onnx、onnxsim、onnxruntime
  • 模型文件就绪:Pipeline 验证阶段 :weights/yolo26n.pt(02 步骤已下载)真实训练阶段 :runs/detect/yolo26_drone_n640/weights/best.pt(09 训练产出)

步骤

3.1 导出命令(Ultralytics 内置)

Pipeline 验证阶段(首次执行,用预训练权重):

bash 复制代码
yolo export \
    model=weights/yolo26n.pt \
    format=onnx \
    imgsz=640 \
    batch=1 \
    opset=11 \
    simplify=True \
    dynamic=False \
    nms=False

真实训练阶段(数据集 + 训练完成后):

bash 复制代码
yolo export \
    model=runs/detect/yolo26_drone_n640/weights/best.pt \
    format=onnx \
    imgsz=640 \
    batch=1 \
    opset=11 \
    simplify=True \
    dynamic=False \
    nms=False

或者直接运行weight目录下的export_onnx.py脚本(也要在yolo26的conda环境中)

python 复制代码
from ultralytics import YOLO

model = YOLO("weights/yolo26n.pt")
model.export(
    format="onnx",
    imgsz=640,
    batch=1,
    opset=11,
    simplify=True,
    dynamic=False,
    nms=False,
)

成功导出

关键参数说明

参数 值 为什么这样设
format onnx 导出为 ONNX 格式
imgsz 640 必须与训练时一致,ATC 要求固定 shape
batch 1 边缘部署通常 batch=1,ATC 转换更简单
opset 11 Atlas 310B1 CANN 23.0.rc3 对 opset 11 兼容性最好
simplify True 去除冗余算子,减少 ATC 不支持风险
dynamic False 必须固定 shape,Atlas 310B1 不支持动态 shape
nms False 不在 ONNX 内做 NMS,由板端 Python 后处理

关于 opset 的特殊说明: Atlas 310B1 的 CANN 23.0.rc3 对算子支持有限。 opset 11 是安全选择。如果 ATC 报 unsupported operator, 可以尝试降到 opset 9,或检查具体不支持的算子名称。
关于 NMS 的特殊说明: Ultralytics 可以导出包含 NMS 的 ONNX(nms=True), 但 Atlas 310B1 的 ATC 可能不支持 NMS 算子。 因此第一轮必须用 nms=False,在板端用 Python 实现 NMS。

导出结果

成功后会在 model= 指定的同目录下生成同名 .onnx:

bash 复制代码
# Pipeline 验证阶段
weights/
├── yolo26n.pt
└── yolo26n.onnx             # ← 目标文件

# 真实训练阶段
runs/detect/yolo26_drone_n640/weights/
├── best.pt
└── best.onnx                # ← 目标文件

记录好产物路径,后面 3.3 / 3.4 的脚本里要用到。

3.2 ONNX 模型检查(onnx checker)

或者直接运行 weights/ 目录下的 check_onnx.py 脚本(也要在 yolo26 的 conda 环境中)

python 复制代码
"""check_onnx.py --- ONNX checker + 输入/输出 shape 打印"""
import onnx

model = onnx.load("weights/yolo26n.onnx")
onnx.checker.check_model(model)
print("✅ ONNX checker passed")

print("\nInputs:")
for inp in model.graph.input:
    shape = [d.dim_value for d in inp.type.tensor_type.shape.dim]
    print(f"  {inp.name}: {shape}")

print("\nOutputs:")
for out in model.graph.output:
    shape = [d.dim_value for d in out.type.tensor_type.shape.dim]
    print(f"  {out.name}: {shape}")

ops = sorted({node.op_type for node in model.graph.node})
print(f"\nOperators ({len(ops)} 种):")
for op in ops:
    print(f"  - {op}")

成功检查

3.3 PC 端推理基准

在导出 ONNX 后,用 onnxruntime 跑一次 PC 端推理作为基准。 把脚本里的 path/to/model.onnx 和 path/to/test_image.jpg 替换成实际路径。

约定 :测试图放在 02 步骤建立的 weights/test_images/ 下(02:147-156)。

bash 复制代码
python -c "
import onnxruntime as ort
import numpy as np
import cv2

# 加载模型
sess = ort.InferenceSession('path/to/model.onnx')
input_name = sess.get_inputs()[0].name
print(f'Input name: {input_name}, shape: {sess.get_inputs()[0].shape}')

# 读取测试图片并预处理
img = cv2.imread('path/to/test_image.jpg')
img_resized = cv2.resize(img, (640, 640))
blob = img_resized[:, :, ::-1].transpose(2, 0, 1).astype(np.float32) / 255.0
blob = np.expand_dims(blob, axis=0)

# 推理
outputs = sess.run(None, {input_name: blob})
print(f'Output shape: {outputs[0].shape}')
print(f'Output dtype: {outputs[0].dtype}')
print(f'Output range: [{outputs[0].min():.4f}, {outputs[0].max():.4f}]')
"

实际示例(Pipeline 验证阶段):

bash 复制代码
python -c "
import onnxruntime as ort
import numpy as np
import cv2

sess = ort.InferenceSession('weights/yolo26n.onnx')
input_name = sess.get_inputs()[0].name
print(f'Input name: {input_name}, shape: {sess.get_inputs()[0].shape}')

img = cv2.imread('weights/test_images/drone_test.jpg')
img_resized = cv2.resize(img, (640, 640))
blob = img_resized[:, :, ::-1].transpose(2, 0, 1).astype(np.float32) / 255.0
blob = np.expand_dims(blob, axis=0)

outputs = sess.run(None, {input_name: blob})
print(f'Output shape: {outputs[0].shape}')
print(f'Output dtype: {outputs[0].dtype}')
print(f'Output range: [{outputs[0].min():.4f}, {outputs[0].max():.4f}]')
"

这个基准数据非常重要:输出 shape:板端 ACL 推理后必须对齐输出范围:板端推理结果应与此接近后续板端精度对比以此为准

把打印出的 Output shape / dtype / range 一并记下来,与 3.3 的输入/输出 shape 对照确认。

3.5 导出产物清单

完成后,需要传到板端的文件:

bash 复制代码
要传输的文件:
✅ model.onnx                            # ONNX 模型(路径取决于阶段)
✅ weights/test_images/*.jpg              # 测试图片(02 已准备)
✅ metadata.json(如果生成了)             # 模型元数据

不要传输的文件:
❌ *.pt                                  # 板端不需要 PyTorch 权重
❌ __pycache__/                           # 缓存
❌ runs/ 目录                             # 训练输出
❌ weights/yolo26n_simplified.onnx       # 简化版可不传,除非 ATC 拒收原版

具体路径:

阶段 ONNX 路径
Pipeline 验证 weights/yolo26n.onnx
真实训练 runs/detect/yolo26_drone_n640/weights/best.onnx

3.6 ATC 兼容性预检查

在传输到板端之前,可以在 PC 端做一个简单的 ATC 兼容性预判:

bash 复制代码
python -c "
import onnx
model = onnx.load('path/to/model.onnx')
ops = set()
for node in model.graph.node:
    ops.add(node.op_type)
print('ONNX operators used:', sorted(ops))
# 对照 CANN 23.0.rc3 的算子支持列表,确认都在支持范围内
"

实际示例(Pipeline 验证阶段):

bash 复制代码
python -c "
import onnx
model = onnx.load('weights/yolo26n.onnx')
ops = set()
for node in model.graph.node:
    ops.add(node.op_type)
print('ONNX operators used:', sorted(ops))
"

常见不兼容算子(CANN 23.0.rc3 可能不支持):NonMaxSuppression → 用 nms=False 导出规避Resize (opset 11+) → 通常支持,但需确认Squeeze/Unsqueeze → 可能需要 opset 调整自定义算子 → 避免使用,Ultralytics 标准导出通常没有

3.7 常见导出问题

问题 原因 解决方案
导出后 shape 不含 batch 维 Ultralytics 版本差异 确认 batch=1 参数生效
onnx checker 报错 模型包含不支持的算子 降低 opset 到 9,或升级 onnxsim
ONNX 文件过大(>500MB) 板端磁盘不够 用 yolo26n 而非 yolo26m/l/x
输出 shape 与 3.3 记录的不一致 简化步骤改变了节点 以最终传给 ATC 的那份为准
simplify 后形状变化 Simplify 改变了节点 重新检查输入输出 shape

注意事项

  • Pipeline 验证阶段和真实训练阶段要分别执行一次本步骤 ,命令几乎一样,只换 model= 参数
  • ONNX 输出 shape 是后续 06 ACL 解码 / 07 端到端验证的关键依赖, 跑完 3.3 / 3.4 一定要把实际 shape 记下来,不要假设
  • 每个阶段的 ONNX 产物路径不一样,后续步骤引用时按上表对号入座
  • 如果 ATC 报算子不支持,优先尝试:降低 opset 到 9用 3.2 的 onnxsim 再简化一次切换到更小的模型(如 yolo26s → yolo26n)
  • 不要在 PC 端用预训练权重 yolo26n.pt 的推理结果做业务判断 ------ 它检测不出无人机,只用来验证 pipeline
相关推荐
胡二拉二胡42 分钟前
抗遗忘单词表 App:让单词记忆不再「背了就忘」
经验分享·笔记·学习
2301_789380491 小时前
【踩坑记录】Windows资源管理器无限重启,改注册表、看日志
windows·经验分享·笔记
by209991 小时前
学会使用 std::string 类,并理解其内部是如何管理字符串的详细阐述(下)
c++·经验分享·笔记·字符串·类和对象·string
DevNo1 小时前
我的2026中医执业医师笔试资料使用小记
笔记·学习·考研
一隅论数智2 小时前
给AI Agent一颗“私域大脑“:本体增强的工程化之路
大数据·运维·数据仓库·人工智能·笔记·学习·政务
zhangrelay2 小时前
ubuntu图形化和命令行查看硬件参数方式汇总-Surface Go
linux·笔记·学习·ubuntu
栈溢出了2 小时前
VS Code Python 自动导包失效排查笔记
开发语言·笔记·python
yi0113 小时前
DAY19: LeetCode 28 找出字符串中第一个匹配项的下标
人工智能·笔记·python·算法·leetcode
zhangrelay14 小时前
机械臂末端规划智能模型的误区演示
linux·笔记·学习·ubuntu·ros2