Yolov5 export.py实现onnx模型的导出

查了很多资料,很多用python代码写的,只需要这个库那个库的,最后都没成功。

不如直接使用Yolov5里面的 export.py实现模型的转换。

一:安装依赖

因为yolov5里面的requirments.txt是将这些转换模型的都注释掉了

所以需要解除注释然后再安装

根据你需要导出模型的类型进行解除注释

然后再安装依赖

复制代码
pip install -r requirements.txt

二:导出模型

根据官方教程来说明如何导出模型,帮助如何将训练好的YOLOv5模型转换成ONNX格式或者TorchScipt格式。

在export.py里面设置模型和数据源的yaml

官方文档写了具体可以导出的类型。

在 --include添加导出的类型。

复制代码
 python export.py  --include torchscript onnx

三:测试和验证推理

复制代码
python detect.py --weights yolov5s.onnx --dnn  
python val.py --weights yolov5s.onnx --dnn  

记住推理的时候,要修改detect.py里面的数据源yaml文件。不然会出现推理的时候标签不一致。

在模型的导出中,onnx不需要GPU进行推理,没有用GPU推理,导致推理的时候很慢,但是tensorRT需要GPU进行推理。

文件有些大,根据官方文档的提示

💡 ProTip: Add --half to export models at FP16 half precision for smaller file sizes

添加--half以FP16半精度导出模型,以缩小文件大小

原来的.pt模型是40.2MB.转成ONNX后模型大小是79.9MB

用完--half完,模型大小变为40.0MB,明显变小了

再次进行推理验证,时间少了43%左右,说明这个是有效果的,测试了一下,对精度影响不大。

相关参数解释:

复制代码
def parse_opt():
    """
    data: 数据集目录 默认=ROOT / 'data/coco128.yaml'
    weights:权重文件目录 默认=ROOT / 'yolov5s.pt'
    img-size: 输入模型的图片size=(height, width) 默认=[640, 640]
    batch-size: batch大小 默认=1
    device: 模型运行设备 cuda device, i.e. 0 or 0,1,2,3 or cpu 默认=cpu
    include: 要将pt文件转为什么格式 可以为单个原始也可以为list 默认=['torchscript', 'onnx', 'coreml']
    half: 是否使用半精度FP16export转换 默认=False
    inplace: 是否设置 YOLOv5 Detect() inplace=True  默认=False
    train: 是否开启model.train() mode 默认=True  coreml转换必须为True
    optimize: TorchScript转化参数 是否进行移动端优化  默认=False
    int8: 支持CoreML/TF INT8 量化 不支持ONNX
    dynamic: ONNX转换参数  dynamic_axes  ONNX转换是否要进行批处理变量  默认=False
    simplify: ONNX转换参数 是否简化onnx模型  默认=False
    opset: ONNX转换参数 设置ONNX版本  默认=13
    topk-per-class: TF.js每一类别都要保留 默认=100
    topk-all: TF.js  Topk为所有class保留
    iou-thres: TF.js IoU threshold   default=0.45
    conf-thres: TF.js  confidence threshold  default=0.25
    include:     需要导出的版本  default=['torchscript', 'onnx'],
    """
复制代码

参考文章:

yolov5 pt 模型 导出 onnx_yolov5 export.py-CSDN博客

TFLite, ONNX, CoreML, TensorRT Export - Ultralytics YOLOv8 Docs

源码解析二 模型转换 export.py - ---dgw博客 - 博客园 (cnblogs.com)

相关推荐
Unknown To Known5 分钟前
血液细胞检测数据集 8类 血小板 淋巴细胞 免疫球蛋白 YOLO格式
yolo
糖炒栗子03267 分钟前
RF-DETR + DeepSpeed 双节点集群训练环境搭建操作记录
python
今天AI了吗27 分钟前
【AI智能体】Hermes Agent 从部署到项目实战操作详解
java·人工智能·python·milvus
心运软件1 小时前
基于机器学习的智能邮件分类系统:从数据采集到模型部署全流程实战
人工智能·python·算法·随机森林·机器学习·分类·scikit-learn
ZHOU_WUYI1 小时前
5. fastwam 模型 video expert pre dit过程
pytorch·python·世界动作模型
guo_xiao_xiao_1 小时前
YOLO[多场景楼梯环境][楼梯]目标检测数据集-24张-barrierfreenavigation_stairs-u437q_v1_24
yolo·目标检测·目标跟踪
过期的秋刀鱼!2 小时前
机器学习开发的迭代循环
人工智能·python·深度学习·神经网络·机器学习
小白学大数据2 小时前
基于Python的抖音网页版视频点赞数增长趋势追踪系统设计与实现
开发语言·爬虫·python·搜索引擎·音视频
向哆哆2 小时前
水稻病害检测数据集分享(适用于YOLO系列深度学习分类检测任务)
深度学习·yolo·分类
circuitsosk2 小时前
平台整体能力与功能特性设计:分群、联运、ABTest与运营位系统
python·机器学习·搜索引擎·ab测试·vllm·rag检索