pytorch训练yolov5s.pt,之后转onnx。

以检测条形码为例,1分类问题,训练yolov5s.pt。

1、github上下载yolov5程序以及pt模型文件,ultralytics/yolov5: YOLOv5 

pt模型训练,转onnx

 pt模型训练,转onnx

2、data文件夹里,images放入样本图片,labels放入对应的txt标签(labelImg制作,注意选择YOLO模式)

pt模型训练,转onnx

 新建yaml文件(上图中的barcode.yaml),用于指定路径和类别。为了省事训练和验证用的一样的图。

pt模型训练,转onnx

3、VSCode打开程序文件夹

 修改models里yolov5s.yaml类别为1

pt模型训练,转onnx

 修改根目录里的train.py

pt模型训练,转onnx

 点击VSCode右上方的开始按钮,训练完毕后,在runs\train\exp\weights\里得到pt模型 best.pt

4、测试下效果

修改根目录里detect.py

pt模型训练,转onnx

 点击VSCode右上方的开始按钮,检测完毕后,在runs\detect\exp\里得到结果图。

5、pt转onnx

VSCode下方的“终端”里输入命令。注意自己的pt路径,会在同路径里得到best.onnx。

python export.py --weights runs/train/exp7/best.pt --img 640 --batch 1

简化模型,先安装包。

【为什么要简化?】在训练完深度学习的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')

 

相关文章:

  • 2023-03-03
  • 2021-05-18
  • 2022-01-08
  • 2022-12-23
  • 2021-11-12
  • 2021-11-02
  • 2022-12-23
  • 2022-12-23
猜你喜欢
  • 2022-12-23
  • 2022-12-23
  • 2022-12-23
  • 2022-12-23
  • 2022-12-23
  • 2021-05-08
  • 2022-12-23
相关资源
相似解决方案