首页yolov5 onnx 推理

yolov5 onnx 推理

时间: 2025-02-10 09:58:38 浏览: 61

### 使用 ONNX Runtime 实现 YOLOv5 模型推理为了提高模型推理速度，将 PyTorch 训练好的 YOLOv5 模型转换为 ONNX 格式并使用 ONNX Runtime 进行推理是一种有效的方法[^1]。 #### 准备工作确保安装必要的库： ```bash pip install onnxruntime numpy opencv-python torch torchvision ``` #### 将 YOLOv5 转换为 ONNX 格式通过官方脚本可以方便地完成这一过程。假设已经有一个训练好的 `yolov5s.pt` 文件： ```python import torch from pathlib import Path model = torch.hub.load('ultralytics/yolov5', 'custom', path='path/to/best.pt') dummy_input = torch.randn(1, 3, 640, 640) torch.onnx.export( model, dummy_input, "yolov5s.onnx", input_names=['input'], output_names=['output'], opset_version=12 ) ``` #### 加载 ONNX 模型并执行推理下面展示了如何加载刚刚保存下来的 `.onnx` 文件，并对其进行图像检测操作： ```python import cv2 import numpy as np import onnxruntime as ort def preprocess(image_path): img = cv2.imread(image_path) img_resized = cv2.resize(img, (640, 640)) blob = cv2.dnn.blobFromImage(img_resized, scalefactor=1/255.0, size=(640, 640), swapRB=True, crop=False) return blob session = ort.InferenceSession("yolov5s.onnx") image_blob = preprocess("test.jpg") # 替换成实际图片路径 inputs = {session.get_inputs()[0].name: image_blob} outputs = session.run(None, inputs)[0] print(outputs.shape) # 输出形状应类似于(batch_size, num_boxes, box_info) ``` 这段代码实现了从读取输入图像到预处理再到最终调用 ONNX Runtime 执行预测的过程。

阅读全文