YOLOv8 模型转换为 ONNX 格式并进行优化与验证

最新推荐文章于 2025-04-02 15:02:42 发布

mr.Darker

最新推荐文章于 2025-04-02 15:02:42 发布

阅读量2.1k

点赞数 21

分类专栏：深度探索文章标签： YOLO

本文链接：https://blog.csdn.net/m0_58648890/article/details/144189685

版权

深度探索专栏收录该内容

5 篇文章

订阅专栏

在深度学习应用中，ONNX（Open Neural Network Exchange）格式被广泛使用，它允许将不同框架（如PyTorch、TensorFlow等）训练的模型进行跨平台使用。在本文中，我们将探讨如何将YOLOv8模型从PyTorch转换为ONNX格式，进行优化，并验证优化后的ONNX模型的有效性。

1. 准备工作

首先，我们需要安装一些必要的依赖包，包括Ultralytics（用于YOLOv8模型的导出）、onnx（用于加载和保存ONNX模型）以及onnxruntime（用于执行和优化ONNX模型）。请确保在Python环境中安装以下依赖项：

pip install ultralytics onnx onnxruntime

2. 转换YOLOv8模型为ONNX格式

在使用YOLOv8模型时，我们通常从PyTorch格式的.pt文件开始。使用Ultralytics库提供的export功能可以将模型轻松转换为ONNX格式。

步骤1：加载YOLOv8模型并导出为ONNX格式

from ultralytics import YOLO

# 加载YOLOv8模型
model = YOLO("yolov8n.pt")

# 导出模型为ONNX格式
model.export(format="onnx")

在这段代码中，我们使用YOLO类加载PyTorch模型（yolov8n.pt），并使用export方法将其导出为ONNX格式。

步骤2：查看导出的ONNX模型

导出的ONNX模型会保存为一个.onnx文件。在本例中，模型将保存为yolov8n.onnx。你可以使用Netron等工具可视化ONNX模型，查看网络结构和层级信息。

3. 优化ONNX模型

ONNX模型通常会有一些冗余操作，在部署到生产环境时需要进行优化。优化后的ONNX模型可以提高推理性能，减少计算和内存消耗。以下是优化过程的步骤。

步骤1：安装`onnxruntime`（如果尚未安装）

onnxruntime是一个用于高效运行ONNX模型的库，也提供了一些优化选项。请确保已经安装onnxruntime。

pip install onnxruntime

步骤2：使用`onnxruntime`进行优化

我们可以使用onnxruntime中的GraphOptimizationLevel来优化ONNX模型。onnxruntime提供了不同的优化级别，分别为：

ORT_ENABLE_BASIC: 基本优化
ORT_ENABLE_EXTENDED: 扩展优化
ORT_ENABLE_ALL: 所有可用优化

优化后的模型将会减少计算量，从而提高推理速度。

import onnx
import onnxruntime as ort
from onnxruntime import GraphOptimizationLevel, SessionOptions

# 加载ONNX模型
onnx_model = onnx.load('yolov8n.onnx')

# 设置优化级别
optimization_level = GraphOptimizationLevel.ORT_ENABLE_BASIC  # 可以选择不同的优化级别
options = SessionOptions()
options.graph_optimization_level = optimization_level

# 使用onnxruntime进行优化
session = ort.InferenceSession('yolov8n.onnx', options)

# 保存优化后的ONNX模型
onnx.save(onnx_model, 'yolov8n_optimized.onnx')

在上面的代码中，我们通过onnxruntime加载ONNX模型，并设置优化级别为ORT_ENABLE_BASIC，然后使用InferenceSession优化模型并将其保存为yolov8n_optimized.onnx。

4. 验证ONNX模型的有效性

在进行模型优化后，确保模型依然能够正确地进行推理非常重要。验证过程包括加载优化后的ONNX模型，并使用一些输入数据进行推理测试。

步骤1：加载优化后的ONNX模型并进行推理

我们可以通过onnxruntime加载优化后的ONNX模型并执行推理，确保模型在推理时不会出现错误。

import onnxruntime as ort
import numpy as np

# 加载优化后的ONNX模型
session = ort.InferenceSession('yolov8n_optimized.onnx')

# 创建一个假设的输入数据，通常是一个图像
# 假设输入数据为形状 (1, 3, 640, 640)，即批次大小为1，3通道，640x640的图像
input_data = np.random.randn(1, 3, 640, 640).astype(np.float32)

# 获取输入层的名称（通常是"images"）
input_name = session.get_inputs()[0].name

# 获取输出层的名称（YOLO模型通常有多个输出）
output_name = session.get_outputs()[0].name

# 执行推理
outputs = session.run([output_name], {input_name: input_data})

# 输出推理结果
print(outputs)

在这段代码中，我们通过onnxruntime加载优化后的ONNX模型，并用随机数据模拟输入图像进行推理。session.run返回的是模型的输出结果。