pytorch onnx 推理 yolov5

原创

mob649e81673fa5 2024-08-16 07:22:25 ©著作权

文章标签 加载模拟输入 python 文章分类 PyTorch 人工智能

©著作权归作者所有：来自51CTO博客作者mob649e81673fa5的原创作品，请联系作者获取转载授权，否则将追究法律责任

如何实现 PyTorch ONNX 推理 YOLOv5

在深度学习项目中，将 PyTorch 模型导出为 ONNX 格式并进行推理是一项常见的需求。YOLOv5 是一种高效的目标检测模型，本文将指导你如何将其推理过程实现PyTorch ONNX 推理 YOLOv5。

整体流程

下面是你需要遵循的步骤流程：

步骤	描述
1	安装必要的库物品，如 PyTorch 和 ONNX
2	下载 YOLOv5 模型
3	训练或加载预训练模型
4	将模型导出为 ONNX 格式
5	编写推理代码，使用 ONNX Runtime 进行推理

journey
    title PyTorch ONNX 推理 YOLOv5 过程
    section 安装必要库
      安装 PyTorch               : 5: John, Sarah
      安装 ONNX                  : 4: Mary, Alex
    section 下载模型
      下载 YOLOv5                : 4: John
    section 加载模型
      训练模型或加载预训练模型 : 3: Sarah
    section 导出模型
      导出为 ONNX 格式         : 4: Alex
    section 推理
      使用 ONNX Runtime 推理    : 5: Mary

详细步骤与代码

1. 安装必要的库

首先，在你的 Python 环境中安装 PyTorch 和 ONNX:

pip install torch torchvision onnx onnxruntime

这条命令将会安装 PyTorch 及其相关库，还有用于处理 ONNX 模型的 onnx 和 onnxruntime。

2. 下载 YOLOv5 模型

接下来，从 GitHub 下载 YOLOv5 模型代码，打开命令行并运行以下命令：

git clone 
cd yolov5
pip install -r requirements.txt

这将会下载 YOLOv5 的代码并安装所需的依赖。

3. 训练或加载预训练模型

如果你想使用预训练模型，可以直接加载 YOLOv5 提供的模型：

import torch

# 加载预训练模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

解释：此代码使用 PyTorch 的 torch.hub 方法从 GitHub 加载 YOLOv5 小模型。

4. 将模型导出为 ONNX 格式

一旦你加载了模型，就可以导出 ONNX 格式的模型：

dummy_input = torch.randn(1, 3, 640, 640)  # 模拟输入
torch.onnx.export(model, dummy_input, "yolov5.onnx", 
                  verbose=False, 
                  input_names=['input'], 
                  output_names=['output'],
                  opset_version=11)  # 指定 ONNX opset 版本

解释：

dummy_input 是模拟输入，表示输入图像的张量形状。
torch.onnx.export 方法用于导出模型，第一个参数是要导出的模型，第二个是模拟输入，第三个是输出的 ONNX 文件名。

5. 使用 ONNX Runtime 进行推理

最后，你可以使用 ONNX Runtime 进行推理。以下是一个简单的推理示例：

import onnxruntime
import numpy as np
from PIL import Image

# 加载 ONNX 模型
session = onnxruntime.InferenceSession("yolov5.onnx")

# 处理输入
img = Image.open("path_to_your_image.jpg").resize((640, 640))
img = np.array(img).astype('float32')  # 转换为浮点数格式
img = img.transpose(2, 0, 1)  # CHW 格式
img = np.expand_dims(img, axis=0)  # 添加 batch size 维度

# 推理
outputs = session.run(["output"], { "input": img })[0]

print(outputs)  # 查看输出结果

解释：

通过 onnxruntime.InferenceSession 加载导出的 ONNX 模型。
使用 PIL 库处理图像并将其转换为适合模型输入的格式。
进行推理并输出结果。

gantt
    title PyTorch ONNX 推理 YOLOv5 时间表
    dateFormat  YYYY-MM-DD
    section 安装步骤
    安装库          :active, 2023-10-01, 1d
    section 下载模型
    下载YOLOv5      :2023-10-02, 1d
    section 加载模型
    加载或训练模型 : 2023-10-03, 2d
    section 导出模型
    导出为ONNX格式 : 2023-10-05, 1d
    section 进行推理
    使用推理        : 2023-10-06, 1d