YOLOv5(ultralytics) pytorch转onnx

最新推荐文章于 2024-08-08 17:02:55 发布

THE@JOKER

最新推荐文章于 2024-08-08 17:02:55 发布

阅读量1.2k

点赞数 1

分类专栏： TensorRT 文章标签： pytorch 深度学习 python

本文链接：https://blog.csdn.net/W1995S/article/details/119335648

版权

TensorRT 专栏收录该内容

7 篇文章 2 订阅

订阅专栏

1.官方自带的export.py，选择模型

python export.py --weights weights/yolov5s.pt --img 640 --batch 1

2.根据错误提示pip install coremltools、packaging。然后继续运行1的命令
3.此时 weights下出现三个文件 :onnx、mlmodel、torchscript
4.使用神经网络Netron,出现网络结构图。

import netron
netron.start('yolov5s.onnx')

5 . pt模型转为了onnx模型，可以说已经脱离了pytorch了，但是转ncnn还需要对模型做一个Simplifier操作，因为转出的onnx模型还有许多冗余，这在ncnn里也是不支持的，避免转ncnn时各种报错安装简化打开weights简化

python -m onnxsim yolov5s.onnx yolov5s_sim.onnx

使用onnx推理

import onnx
import cv2
from utils import datasets,torch_utils
import numpy as np
import torch
import onnxruntime
from utils.general import (
    check_img_size, non_max_suppression, apply_classifier, scale_coords,
    xyxy2xywh,  strip_optimizer, set_logging)

src_img = cv2.imread("data/images/bus.jpg")

img = src_img[:, :, ::-1] #BGR TO RGB
img = datasets.letterbox(img, new_shape=(640,640))[0]
img = cv2.resize(img,dsize=(640,640),interpolation=cv2.INTER_NEAREST)
img = img.transpose(2, 0, 1)
img = img.astype(np.float32)
img /= 255. # (3, 640, 640)


ort_session = onnxruntime.InferenceSession("weights/yolov5s-sim.onnx")
pred = ort_session.run(None, {ort_session.get_inputs()[0].name: img[None]})[0]


det = non_max_suppression(torch.tensor(pred), 0.3, 0.5, classes=None, agnostic=False)[0]
print(det)
if det is not None and len(det):
    det[:, :4] = scale_coords(img.shape[1:], det[:, :4], src_img.shape).round()
    boxes = det.cpu().detach().numpy()
    for box in boxes:
        box = box.astype(np.int)
        cv2.rectangle(src_img, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)

cv2.imshow("....", src_img)
cv2.waitKey(0)
cv2.destroyAllWindows()

https://www.csdn.net/tags/NtjaAg5sNjk1NDYtYmxvZwO0O0OO0O0O.html
https://blog.csdn.net/yx868yx/article/details/116663020?utm_medium=distribute.pc_aggpage_search_result.none-task-blog-2~aggregatepage~first_rank_v2~rank_aggregation-4-116663020.pc_agg_rank_aggregation&utm_term=yolov5%E8%BD%AC%E6%8D%A2onnx&spm=1000.2123.3001.4430
https://blog.csdn.net/qq_45057749/article/details/115016683?spm=1001.2014.3001.5501
https://blog.csdn.net/qq_45057749/article/details/115013509?spm=1001.2014.3001.5501
https://blog.csdn.net/weixin_41868104/article/details/109359343
https://blog.csdn.net/qq_33160678/category_10419575.html