pt模型训练,转onnx
pytorch训练yolov5s.pt,之后转onnx。
以检测条形码为例,1分类问题,训练yolov5s.pt。
1、github上下载yolov5程序以及pt模型文件,https://github.com/ultralytics/yolov5/releases
2、data文件夹里,images放入样本图片,labels放入对应的txt标签(labelImg制作,注意选择YOLO模式)
新建yaml文件(上图中的barcode.yaml),用于指定路径和类别。为了省事训练和验证用的一样的图。
3、VSCode打开程序文件夹
修改models里yolov5s.yaml类别为1
修改根目录里的train.py
点击VSCode右上方的开始按钮,训练完毕后,在runs\train\exp\weights\里得到pt模型 best.pt
4、测试下效果
修改根目录里detect.py
点击VSCode右上方的开始按钮,检测完毕后,在runs\detect\exp\里得到结果图。
5、pt转onnx
VSCode下方的“终端”里输入命令。注意自己的pt路径,会在同路径里得到best.onnx。
python export.py --weights runs/train/exp7/best.pt --include onnx --opset 12 --dynamic
简化模型,先安装包。
【为什么要简化?】在训练完深度学习的pytorch或者tensorflow模型后,有时候需要把模型转成 onnx,但是很多时候,很多节点比如cast节点,Identity 这些节点可能都不需要,我们需要进行简化,这样会方便我们把模型转成ncnn或者mnn等这些端侧部署的模型格式或者通过tensorRT进行部署。
pip install onnx coremltools onnx-simplifier
进入best.onnx目录
cd runs/train/exp7/weights
简化,同路径下得到简化后的模型 best-sim.onnx
python -m onnxsim best.onnx best-sim.onnx
查看网络结构
pip install netron python import netron netron.start('best-sim.onnx')
测试下效果
1 | python detect.py - - weights best - sim.onnx - - source path / to / img.jpg |
【代码里参数讲解】
https://blog.csdn.net/IT_charge/article/details/119208680?spm=1001.2101.3001.6650.1&depth_1-utm_relevant_index=2
【推荐】国内首个AI IDE,深度理解中文开发场景,立即下载体验Trae
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· 震惊!C++程序真的从main开始吗?99%的程序员都答错了
· 【硬核科普】Trae如何「偷看」你的代码?零基础破解AI编程运行原理
· 单元测试从入门到精通
· 上周热点回顾(3.3-3.9)
· winform 绘制太阳,地球,月球 运作规律
2021-01-26 Qt 鼠标坐标