将 YOLOv10 模型从 PyTorch 转换为 ONNX

在深度学习模型的部署过程中，将模型从一种格式转换为另一种格式是非常常见的操作。本文将详细介绍如何将 YOLOv10 模型从 PyTorch 的 .pt 格式转换为 ONNX 的 .onnx 格式。

1. 环境准备

首先，我们需要确保环境已经配置好，并且安装了必要的依赖库。以下是环境准备步骤：

1.1 克隆 YOLOv5 仓库

git clone https://github.com/ultralytics/yolov5
cd yolov5

1.2 安装依赖库

使用 pip 安装 YOLOv5 所需的依赖库：

pip install -r requirements.txt

2. 转换模型

在环境准备完成后，我们可以开始将 YOLOv5 模型从 .pt 格式转换为 .onnx 格式。

2.1 使用 YOLOv5 提供的命令进行转换

YOLOv5 提供了一个方便的命令行工具来执行模型转换。假设我们有一个名为 yolov10x_best.pt 的模型文件，我们可以使用以下命令将其转换为 ONNX 格式：

yolo export model=../learn-yolo/models/yolov10x_best.pt format=onnx

2.2 转换过程

执行上述命令后，YOLOv5 会自动将模型转换为 ONNX 格式，并输出转换后的模型文件路径。以下是转换过程的日志输出：

Ultralytics YOLOv8.2.93 🚀 Python-3.12.3 torch-2.4.1+cu121 CPU (13th Gen Intel Core(TM) i9-13900KF)
YOLOv10x summary (fused): 503 layers, 31,605,266 parameters, 0 gradients, 169.9 GFLOPsPyTorch: starting from '../learn-yolo/models/yolov10x_best.pt' with input shape (1, 3, 640, 640) BCHW and output shape(s) (1, 300, 6) (61.2 MB)
requirements: Ultralytics requirement ['onnx>=1.12.0'] not found, attempting AutoUpdate...
WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
Collecting onnx>=1.12.0Downloading onnx-1.16.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (16 kB)
Requirement already satisfied: numpy>=1.20 in /root/miniconda3/lib/python3.12/site-packages (from onnx>=1.12.0) (1.26.4)
Collecting protobuf>=3.20.2 (from onnx>=1.12.0)Downloading protobuf-5.28.1-cp38-abi3-manylinux2014_x86_64.whl.metadata (592 bytes)
Downloading onnx-1.16.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (15.9 MB)━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 15.9/15.9 MB 3.3 MB/s eta 0:00:00
Downloading protobuf-5.28.1-cp38-abi3-manylinux2014_x86_64.whl (316 kB)━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 316.6/316.6 kB 164.3 MB/s eta 0:00:00
Installing collected packages: protobuf, onnx
Successfully installed onnx-1.16.2 protobuf-5.28.1Results saved to /mnt/e/workspace/learn-yolo/models
Predict:         yolo predict task=detect model=../learn-yolo/models/yolov10x_best.onnx imgsz=640
Validate:        yolo val task=detect model=../learn-yolo/models/yolov10x_best.onnx imgsz=640 data=./doclaynet.yaml
Visualize:       https://netron.app
💡 Learn more at https://docs.ultralytics.com/modes/export

2.3 转换结果

转换完成后，模型文件 yolov10x_best.onnx 将保存在指定的目录中。你可以使用以下命令查看转换后的模型文件：

ls ../learn-yolo/models/yolov10x_best.onnx

3. 验证转换后的模型

为了确保转换后的模型能够正常工作，你可以使用以下命令进行预测和验证：

3.1 预测

yolo predict task=detect model=../learn-yolo/models/yolov10x_best.onnx imgsz=640

3.2 验证

yolo val task=detect model=../learn-yolo/models/yolov10x_best.onnx imgsz=640 data=./doclaynet.yaml

4. 总结

通过以上步骤，我们成功地将 YOLOv5 模型从 PyTorch 的 .pt 格式转换为 ONNX 的 .onnx 格式。ONNX 格式的模型可以在多种深度学习框架和硬件平台上运行，方便了模型的部署和推理。

希望本文对你有所帮助！如果你有任何问题或建议，欢迎在评论区留言。

---

参考链接：

• YOLOv5 GitHub 仓库
• ONNX 官方文档
• Ultralytics YOLOv5 文档