使用的yolov5代码:https://github.com/ultralytics/yolov5/tree/v5.0
1.安装依赖库
pip install onnx coremltools onnx-simplifier
2.导出 onnx
python models/export.py --weights yolov5s-best.pt --img 640 --batch 1
在 yolov5s-best.pt 的同级目录生成了 yolov5s-best.onnx、yolov5s-best.mlmodel 和 yolov5s-best.torchscript.pt三个文件。
可以使用netron查看yolov5s-best.onnx模型结构。
# 在终端运行
pip install netron
python
import netron
netron.start('yolov5s-best.onnx')
3.用 onnx-simplifer 简化模型
【为什么要简化?】在训练完深度学习的pytorch或者tensorflow模型后,有时候需要把模型转成 onnx,但是很多时候,很多节点比如cast节点,Identity 这些节点可能都不需要,我们需要进行简化,这样会方便我们把模型转成ncnn或者mnn等这些端侧部署的模型格式或者通过tensorRT进行部署。
python -m onnxsim yolov5s-best.onnx yolov5s-best-sim.onnx
前者yolov5s-best.onnx是需要简化的onnx的路径,后者yolov5s-best-sim.onnx是输出简化后模型的路径。
要继续转成ncnn,请参考yolov5模型转换(二) onnx转ncnn.
参考文献:
YOLOv5利用ncnn部署系列(二).