Python 官方文档:入门教程 => 点击学习
博主想拿官网的yolov5训练好pt模型,然后转换成rknn模型,然后在瑞芯微开发板上调用模型检测。但是官网的版本对npu不友好,所以采用改进结构的版本: 将Focus层改成Conv层将Swis
博主想拿官网的yolov5训练好pt模型,然后转换成rknn模型,然后在瑞芯微开发板上调用模型检测。但是官网的版本对npu不友好,所以采用改进结构的版本:
- 将Focus层改成Conv层
- 将Swish激活函数改成Relu激活函数
自带的预训练模型是预测80类CoCo数据集的yolov5s改进结构,下面就带大家一起转换模型!
首先部署好yolov5的环境,保证可以运行detect.py进行检测,将自己训练好的pt模型放到weights目录下,我这里命名是best.pt。
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 1Namespace(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$
导出的onnx模型可以用Netron网站查看网络结构:
可以看到转换后的结构比较乱,有些结点其实可以跳过,减小模型。这里可以使用修改onnx模型的工具,说明都在Readme中。
这里我们使用深度学习的库onnx-simplifier,通过pip安装即可。
Collecting onnx-simplifier Downloading 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:00Requirement already satisfied: onnx in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnx-simplifier) (1.12.0)Collecting rich Downloading rich-12.6.0-py3-none-any.whl (237 kB) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 237.5/237.5 kB 191.1 kB/s eta 0:00:00Requirement 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.0 Downloading commonmark-0.9.1-py2.py3-none-any.whl (51 kB) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 51.1/51.1 kB 215.0 kB/s eta 0:00:00Collecting pygments<3.0.0,>=2.6.0 Downloading Pygments-2.13.0-py3-none-any.whl (1.1 MB) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.1/1.1 MB 142.8 kB/s eta 0:00:00Installing collected packages: commonmark, pygments, rich, onnx-simplifierSuccessfully 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.onnxInstalling onnxruntime by `/home/dzh/anaconda3/envs/yolov5/bin/python -m pip install --user onnxruntime`, please wait for a moment..Collecting onnxruntime Downloading 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:00Collecting coloredlogs Downloading coloredlogs-15.0.1-py2.py3-none-any.whl (46 kB) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 46.0/46.0 kB 74.5 kB/s eta 0:00:00Requirement already satisfied: protobuf in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnxruntime) (3.19.6)Collecting sympy Downloading sympy-1.11.1-py3-none-any.whl (6.5 MB) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6.5/6.5 MB 72.3 kB/s eta 0:00:00Requirement 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 flatbuffers Downloading flatbuffers-22.9.24-py2.py3-none-any.whl (26 kB)Collecting humanfriendly>=9.1 Downloading humanfriendly-10.0-py2.py3-none-any.whl (86 kB) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 86.8/86.8 kB 129.3 kB/s eta 0:00:00Requirement 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.19 Downloading mpmath-1.2.1-py3-none-any.whl (532 kB) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 532.6/532.6 kB 128.3 kB/s eta 0:00:00Installing 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个通道的图像,所以在图片输入之前改成模型需要的通道数 |
来源地址:https://blog.csdn.net/qq_42257666/article/details/127244265
--结束END--
本文标题: 【yolov5】pytorch模型导出为onnx模型
本文链接: https://www.lsjlt.com/news/407305.html(转载时请注明来源链接)
有问题或投稿请发送至: 邮箱/279061341@qq.com QQ/279061341
下载Word文档到电脑,方便收藏和打印~
2024-03-01
2024-03-01
2024-03-01
2024-02-29
2024-02-29
2024-02-29
2024-02-29
2024-02-29
2024-02-29
2024-02-29
回答
回答
回答
回答
回答
回答
回答
回答
回答
回答
0