【yolov5】pytorch模型导出为onnx模型

博主想拿官网的yolov5训练好pt模型，然后转换成rknn模型，然后在瑞芯微开发板上调用模型检测。但是官网的版本对npu不友好，所以采用改进结构的版本：

将Focus层改成Conv层

将Swish激活函数改成Relu激活函数

自带的预训练模型是预测80类CoCo数据集的yolov5s改进结构，下面就带大家一起转换模型！

1、首先部署好yolov5的环境，保证可以运行detect.py进行检测，将自己训练好的pt模型放到weights目录下，我这里命名是best.pt。

2、pip install onnx安装好onnx库

3、输入下面命令导出模型（后面的weights、img和batch参数可以不加，设置下默认参数即可）

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

然后会分解网络层，预览模型层数和参数量，最后保存为./weights/best.onnx。

(yolov5) dzh@dzh-Lenovo-Legion-Y7000:~/airockchip-yolov5$ python models/export.py --weights ./weights/best.pt --img 640 --batch 1
Namespace(batch_size=1, img_size=[640, 640], weights='./weights/best.pt')
Fusing layers...
Model Summary: 140 layers, 7.2627e+06 parameters, 0 gradientsStarting TorchScript export with torch 1.12.1+cu102...
/home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages/torch/jit/_trace.py:967: TracerWarning: Encountering a list at the output of the tracer might cause the trace to be incorrect, this is only valid if the container structure does not change based on the module's inputs. Consider using a constant container instead (e.g. for `list`, use a `tuple` instead. for `dict`, use a `NamedTuple` instead). If you absolutely need this and know the side effects, pass strict=False to trace() to allow this behavior.module._c._create_method_from_trace(
TorchScript export success, saved as ./weights/best.torchscript.ptStarting ONNX export with onnx 1.12.0...
/home/dzh/airockchip-yolov5/./models/yolo.py:103: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!if augment:
/home/dzh/airockchip-yolov5/./models/yolo.py:128: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!if profile:
/home/dzh/airockchip-yolov5/./models/yolo.py:143: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!if profile:
/home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages/torch/onnx/symbolic_helper.py:621: UserWarning: You are trying to export the model with onnx:Resize for ONNX opset version 10. This operator might cause results to not match the expected results by PyTorch.
ONNX's Upsample/Resize operator did not match Pytorch's Interpolation until opset 11. Attributes to determine how to transform the input were added in onnx:Resize in opset 11 to support Pytorch's behavior (like coordinate_transformation_mode and nearest_mode).
We recommend using opset 11 and above for models using this operator.warnings.warn(
ONNX export success, saved as ./weights/best.onnxExport complete (4.98s). Visualize with https://github.com/lutzroeder/netron.
(yolov5) dzh@dzh-Lenovo-Legion-Y7000:~/airockchip-yolov5$

4、导出的onnx模型可以用Netron网站查看网络结构：

可以看到转换后的结构比较乱，有些结点其实可以跳过，减小模型。这里可以使用修改onnx模型的工具，说明都在Readme中。

这里我们使用深度学习的库onnx-simplifier，通过pip安装即可。

Collecting onnx-simplifierDownloading onnx_simplifier-0.4.8-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (2.0 MB)━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.0/2.0 MB 203.7 kB/s eta 0:00:00
Requirement already satisfied: onnx in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnx-simplifier) (1.12.0)
Collecting richDownloading rich-12.6.0-py3-none-any.whl (237 kB)━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 237.5/237.5 kB 191.1 kB/s eta 0:00:00
Requirement already satisfied: numpy>=1.16.6 in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnx->onnx-simplifier) (1.23.3)
Requirement already satisfied: protobuf<=3.20.1,>=3.12.2 in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnx->onnx-simplifier) (3.19.6)
Requirement already satisfied: typing-extensions>=3.6.2.1 in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnx->onnx-simplifier) (4.4.0)
Collecting commonmark<0.10.0,>=0.9.0Downloading commonmark-0.9.1-py2.py3-none-any.whl (51 kB)━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 51.1/51.1 kB 215.0 kB/s eta 0:00:00
Collecting pygments<3.0.0,>=2.6.0Downloading Pygments-2.13.0-py3-none-any.whl (1.1 MB)━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.1/1.1 MB 142.8 kB/s eta 0:00:00
Installing collected packages: commonmark, pygments, rich, onnx-simplifier
Successfully installed commonmark-0.9.1 onnx-simplifier-0.4.8 pygments-2.13.0 rich-12.6.0
(yolov5) dzh@dzh-Lenovo-Legion-Y7000:~/airockchip-yolov5$ python -m onnxsim ./weights/red.onnx ./weights/red2.onnx
Installing onnxruntime by `/home/dzh/anaconda3/envs/yolov5/bin/python -m pip install --user onnxruntime`,
please wait for a moment..
Collecting onnxruntimeDownloading onnxruntime-1.12.1-cp38-cp38-manylinux_2_27_x86_64.whl (4.9 MB)━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 4.9/4.9 MB 52.8 kB/s eta 0:00:00
Collecting coloredlogsDownloading coloredlogs-15.0.1-py2.py3-none-any.whl (46 kB)━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 46.0/46.0 kB 74.5 kB/s eta 0:00:00
Requirement already satisfied: protobuf in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnxruntime) (3.19.6)
Collecting sympyDownloading sympy-1.11.1-py3-none-any.whl (6.5 MB)━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6.5/6.5 MB 72.3 kB/s eta 0:00:00
Requirement already satisfied: packaging in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnxruntime) (21.3)
Requirement already satisfied: numpy>=1.21.0 in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnxruntime) (1.23.3)
Collecting flatbuffersDownloading flatbuffers-22.9.24-py2.py3-none-any.whl (26 kB)
Collecting humanfriendly>=9.1Downloading humanfriendly-10.0-py2.py3-none-any.whl (86 kB)━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 86.8/86.8 kB 129.3 kB/s eta 0:00:00
Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from packaging->onnxruntime) (3.0.9)
Collecting mpmath>=0.19Downloading mpmath-1.2.1-py3-none-any.whl (532 kB)━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 532.6/532.6 kB 128.3 kB/s eta 0:00:00
Installing collected packages: mpmath, flatbuffers, sympy, humanfriendly, coloredlogs, onnxruntime

然后输入python -m onnxsim ./weights/red.onnx ./weights/red2.onnx优化onnx模型。可以看到优化前后的表。

总结可能遇到的问题：

ONNX export failure: No module named ‘onnx’
原因：没有安装onnx库导致找不到该模块

RuntimeError: Given groups=1, weight of size [32, 12, 3, 3], expected input[1, 3, 640, 640] to have 12 channels, but got 3 channels instead
模型预期的图像通道是12，但是现在输入是3个通道的图像，所以在图片输入之前改成模型需要的通道数

【yolov5】pytorch模型导出为onnx模型相关推荐

pytorch模型(.pt)转onnx模型(.onnx)的方法详解(1)
1. pytorch模型转换到onnx模型 2.运行onnx模型 3.比对onnx模型和pytorch模型的输出结果我这里重点是第一点和第二点,第三部分比较容易首先你要安装依赖库:onnx ...
ML.NET 示例：将ML.NET模型导出到ONNX
在这个示例中,您将看到如何使用ML.NET来训练回归模型,然后将该模型转换为ONNX格式. 问题开放式神经网络交换即ONNX是一种表示深度学习模型的开放格式.使用ONNX,开发人员可以在最先进的工具 ...
darknet cpp weights模型转换成ONNX模型
整理不易,如果觉得有用,记得点赞收藏和分享哦 1. 下载转换需要的代码文件在下面地址下载代码文件 https://gitee.com/liangjiaxi2019/pytorch-YOLOv4 2. ...
建造属于你的无人驾驶车——（三）Solidworks模型导出为URDF模型
建造属于你的无人驾驶车! 本专栏持续更新中- 程序源码:https://github.com/kkmd66/ZZX_RUN Solidworks模型文件:https://github.com/kkmd ...
Pytorch模型(.pth)转onnx模型(.onnx)
简介 Open Neural Network Exchange(ONNX,开放神经网络交换)格式,是一个用于表示深度学习模型的标准,可使模型在不同框架之间进行转移. ONNX是一种针对机器学习所设计的 ...
4模型导出_项目模型规范总结游戏模型制作的注意事项
点击上方"3D天工坊"关注我本期我们来总结一下模型规范导出等~ 1. 单位,比例统一在建模型前先设置好单位,在同一场景中会用到的模型的单位设置必须一样,模型与模型之间的比例要正确 ...
Paddle2ONNX最新升级：飞桨模型全面支持ONNX协议啦！
点击左上方蓝字关注我们 16年前的电影<功夫>可以说是家喻户晓,并将武侠片推向了一个新的高潮.在电影中,街头小混混阿星偶然间被打通了任督二脉,因善开悟,再有"如来神掌" ...
Pytorch版本MobileNetV3转ONNX然后转om模型使用Pyacl离线推理
Pytorch版本MobileNetV3转ONNX然后转om模型使用Pyacl离线推理概述:本文主要讲述把MobileNet转成华为Altas服务器离线推理om模型的过程,本人在转换过程中也遇到过比 ...
将训练好的pytorch模型的pth文件转换成onnx模型（亲测成功）
将训练好的pytorch模型的pth文件转换成onnx模型(亲测成功) 模型转换声明:本文原创,未经许可严禁转载,原文地址https://blog.csdn.net/hutao1030813002/ ...

【yolov5】pytorch模型导出为onnx模型

【yolov5】pytorch模型导出为onnx模型相关推荐

最新文章

热门文章