YOLOv5模型导出

一、准备

Python>=3.8并且PyTorch==1.7

git clone https://github.com/ultralytics/yolov5  # clone repo
cd yolov5
pip install -r requirements.txt  # base requirements
pip install coremltools>=4.1 onnx>=1.9.0 scikit-learn==0.19.2  # export requirements

二、导出训练模型

以下命令可以将yolov5模型导出为ONNX，TorchScript格式，这里我们以yolov5s.pt模型为例。

# export.py文件在yolov5目录下
python export.py --weights yolov5s.pt --img 640 --batch 1  # export at 640x640 with batch size 1

输出如下：

导出的模型将保存在当前目录，如下截图：