概览硬件软件云平台
OAK 应用
DepthAI
AI
软件栈
故障排除

本页目录

  • 概述
  • 转换
  • 从 PyTorch
  • 从 TensorFlow
  • 验证
  • 输入/输出张量
  • 性能

转换为 ONNX

概述

开放神经网络交换格式（ONNX） 是一种广泛使用的机器学习模型表示格式。它是开源的，并得到众多框架和工具的支持。在此我们收集了一些最常见的将模型转换为 .onnx 格式的方法。建议进行此转换，因为它能为后续在相应 RVC 平台 上的转换提供最多的选项。

转换

从 PyTorch

你可以使用 PyTorch ONNX API 来转换和导出你的模型：
Python
1import torch
2import torch.nn as nn
3
4# 定义模型架构（必须与保存的模型匹配）
5class Model(nn.Module):
6    ...
7    
8# 初始化模型并加载训练好的权重
9model = Model()
10model.load_state_dict(torch.load("model_name.pt"))
11model.eval()  # 将模型设置为推理模式
12
13# 定义输入形状并创建虚拟输入张量
14input_shape = ... # 例如 (1, 3, 512, 288)
15dummy_input = torch.randn(input_shape)
16
17# 将模型导出为 ONNX 格式
18torch.onnx.export(model, dummy_input, "model_name.onnx")
这种转换方式不会改变模型的架构，仅改变模型保存的格式。因此，转换后的模型性能与原始模型保持一致。更多信息请参考官方 PyTorch 导出教程。

YOLO 模型

对于 YOLO 模型的转换，我们推荐使用我们的 tools-cli 包，这是一个专门用于简化转换过程的实用工具。当前支持的系列包括 YOLOv5、YOLOv6、YOLOv7、YOLOv8、YOLOv9、YOLOv10、YOLO11、YOLO12、YOLO26、YOLOE 和 Gold-YOLO。有关确切支持的 YOLO 系列、变体以及当前限制，请参阅上游的 支持模型表。与普通的 PyTorch 转换不同，此过程会修改模型架构以标准化其输出。这使得可以使用原生的 DepthAI DetectionNetwork 节点（或用于姿态和实例分割模型的 YOLOExtendedParser 解析器节点）开箱即用地进行解析。因此，后处理可以完全在设备上完成，消除了主机与设备之间的数据传输，降低了整体延迟。虽然这是推荐的方法，但你仍然可以使用通用转换方法导出 YOLO 模型。但请注意，你需要手动处理输出解析，因为不提供原生支持。此外，请注意，导出的 ONNX 模型可能由于转换器和设备上支持的操作限制而无法直接转换或运行。你可以按如下方式安装 tools-cli 包：
Command Line
1git clone --recursive https://github.com/luxonis/tools.git
2cd tools
3pip install .
安装后，你可以直接从命令行运行转换（确保你在 tools 目录的根目录下运行）：
Command Line
1tools <MODEL>.pt --imgsz "<WIDTH> <HEIGHT>"
你可以在 tools-cli 文档 中找到关于转换参数的更多信息。例如，要转换一个期望输入分辨率为 (512, 288) 的 BGR 图像的 YOLOv6n 模型，请运行以下命令：
Command Line
1tools yolov6n.pt --imgsz "512 288" --encoding BGR
生成的 ONNX 模型输入形状为 (1, 3, 512, 288)，输出形状为 (1,85,36,64)、(1,85,18,32) 和 (1,85,9,16)（与通用转换获得的原始模型输出形状 (1, 2304, 85)、(1, 576, 85) 和 (1, 144, 85) 相反；如果连接成一个张量，则为 (1, 3024, 85)）。

从 TensorFlow

对于 TensorFlow（.pb）、Keras（.h5）、tensorflow.js（.json 和 .bin）或 TensorFlow Lite（.tflite）格式的模型，我们推荐使用 tf2onnx 转换工具。
  • 首先，安装 tensorflow 和 tf2onnx 包：
Command Line
1pip install tensorflow
2pip install -U tf2onnx
  • 其次，可以通过命令行进行转换：
Command Line
1python -m tf2onnx.convert --saved-model tensorflow-model-path --output model.onnx

验证

输入/输出张量

在将 ONNX 转换为 RVC 平台之前，请确保 ONNX 模型的输入/输出张量满足以下条件：
  • 正确定义（从相应层接收输入/向相应层发送输出）。
    • 如果不是，请使用 onnx-modifier 工具重新定义输入/输出张量。
  • 形状为 NCHW 形式（批次大小、颜色通道数、高度、宽度）。
    • 如果不是，请在运行 tf2onnx 工具时定义 --inputs-as-nchw data 和 --outputs-as-nchw 标志。
  • 没有维度是动态的（例如批次大小是固定的）。
    • 如果是动态的，请使用 onnxruntime 库将形状固定（有关详细信息，请参见本教程）。

性能

将模型转换为 ONNX 有时会导致模型性能略有差异。这些差异通常是由于舍入、数值精度、层实现和操作近似值的差异造成的。建议使用各种测试输入来比较原始模型和 ONNX 模型的输出。你可以使用平均绝对误差 (MAE) 或均方误差 (MSE) 等指标来评估差异。如果发现差异，请考虑以下步骤：
  • 更新到最新的 ONNX Opset：确保你使用的是最新的 ONNX opset 版本。较新的 opset 通常包含增强模型兼容性的改进和错误修复。
  • 简化模型操作：在将模型转换为 ONNX 之前，先简化模型中的复杂操作或层。这有助于减少与转换相关的问题。
  • 确保精度一致性：在两个模型中为权重、输入和输出使用相同的数据类型（例如 float32），以保持数值精度。
通过遵循这些步骤，你可以最大限度地减少差异，确保你的 ONNX 模型与原始模型的性能紧密匹配。