Phase A · Step 3：模型导出 ONNX

用户1403605819383
0 阅读7分钟

目标

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

阶段定位

本步骤在 Phase A 的位置：

[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 验证阶段（首次执行，用预训练权重）：

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

真实训练阶段（数据集 + 训练完成后）：

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环境中）

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,
)

成功导出

alt text

alt text

关键参数说明

参数值为什么这样设
formatonnx导出为 ONNX 格式
imgsz640必须与训练时一致，ATC 要求固定 shape
batch1边缘部署通常 batch=1，ATC 转换更简单
opset11Atlas 310B1 CANN 23.0.rc3 对 opset 11 兼容性最好
simplifyTrue去除冗余算子，减少 ATC 不支持风险
dynamicFalse必须固定 shape，Atlas 310B1 不支持动态 shape
nmsFalse不在 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：

# 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 环境中）

"""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}")

成功检查

alt text

3.3 PC 端推理基准

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

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

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 验证阶段）：

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 导出产物清单

完成后，需要传到板端的文件：

要传输的文件：
✅ 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 兼容性预判：

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 验证阶段）：

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
avatar
文章
阅读
粉丝