如何在windows系统下将yolov5的pt模型导出为onnx模型

前言

最近在做本科毕业设计要求是在树莓派上部署yolo算法来实现火灾检测在网上查了很多资料最后选择用yolov5s模型先试着在树莓派上部署看下效果如何由于从大佬那里拿到了yolov5火灾检测模型但想要将它移植到树莓派上第一步要把pt模型转换成onnx模型原因我想大概是因为树莓派算力不如pc因此要转换成合适的模型进行部署。一开始我觉得转换模型应该是很简单的一个事情但是在进行的过程中踩了很多坑走了很多弯路因此在这里总结一下也希望有相同需求的朋友能从这里获得帮助ps:因为本人水平有限所以能帮到大家很高兴。

第一个问题

在进行模型转化时从网上找的代码是这样

python export.py --weights ‘weights/last.pt’ --batch-size 1 --img_size 320

大概内容是调用export.py这个文件来进行模型的转化但是将代码复制到项目命令行里运行后产生如下报错

我看了下报错的原因大概是语法错误后来在网上查找后发现其实这个命令可以不带后面的–batch-size和–img-size如果要带的话这里的–img_size应该改为–img-size。

第二个问题

将后面的参数删掉以后发现还是不能正常运行后来经过分析后发现是因为这里的’weights/last.pt’应该是你项目里这个.pt模型的相对地址因此将代买修改为以下形式

python export.py --weights ‘runs/train/exp_yolov5s/weights/best.pt’

运行后发现还是不行报错如下

在网上查找详细步骤后发现是由于之前并没有安装好onnx的包导致于是又输入了这样一行代码安装缺少的onnx包

pip install onnx

很快就安装好了

安装好了以后运行还是和上面一样的报错于是在网上查找资料发现是因为pytorch的版本太高导致包不兼容所以又重新安装了1.9.0版本的pytorch因为之前安装了conda虚拟环境所以安装过程很顺利不过因为我家网速原因下载速度有些慢…

conda install pytorch1.9.0 torchvision0.10.0 torchaudio==0.9.0 cudatoolkit=10.2 -c pytorch

第三个问题

安装好之后又运行了一编export的文件结果还是报错不过所幸这次报错的原因不一样了这也说明之前那个问题已经解决了虽然不清楚里面的原理不过对于我这个小白来说能用就已经是一件令人激动的事情了如果有大佬的话希望可以帮忙分析一下。这次的报错是这样的

这个问题我在网上搜了很多大部分都是说是conda的环境变量配的有问题然而我的conda早就安装好了而且可以正常运行之前也配置过环境变量所以应该不是这个原因后来偶然间看到有人说可能是因为pillow版本的问题于是我卸载然后重新装了pillow:

pip uninstall pillow

pip install pillow==8.4

第四个问题

经过上面踩过许多次坑之后这次终于它提示成功运行了

但是我还没来得及高兴又来了当头一棒我发现在我存放.pt模型的文件夹里面并没有成功出现导出的onnx文件后来我发现在上面

运行了很多次都是相同的结果然后当时我心态就快崩溃了后来又重振旗鼓去网上找解决方法然后找到了一个看起来很靠谱的博客按照大佬的办法修改了export.py的代码

在这里修改data和weight的默认路径为你数据源data.yaml的路径和模型的权重文件里的best.pt的路径在修改之后在命令行输入这行代码

python export.py --include torchscript onnx

终于程序正常运行我在模型best.pt所在文件夹内看到了导出的.onnx文件和.torchscript文件如果只想导出onnx文件可以把前面的去掉直接输入

python export.py --include onnx

到此将模型部署到树莓派上的第一步总算完成了。

总结

这是本人第一次在csdn上尝试写这种总结类的博客目的也很简单想把这个过程中自己遇到的坑和尝试的经历记录下来方便后面开发过程中有需要的时候回来翻看同时发布到csdn上也希望能够得到大佬们的批评指正如果某些内容能够帮到大家那再好不过了。
作为一个机械专业的大四学生第一次做毕设就拿到这么有挑战性的题目其实对我来说是很恐惧的因为自己之前并没有识别算法这方面的基础知识储备研究生想做的方向是嵌入式不过因为做毕设需要用到深度学习图像识别所以尝试将它部署到树莓派上现在工作才刚刚开始就遇到这么多的困难其实是超出我的预期的不过经历了这些之后我对于开发应用这类工作有了新的认识其实有时候这个东西真的没有你想象中的那么难只是你要做到心平气和不要浮躁多学多问多尝试总会有新的收获加油