yolov5训练模型导出为TorchScript、ONNX、CoreML

在你开始前

克隆这个 repo 并安装requirements.txt依赖项，包括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 模型

此命令将预训练的 YOLOv5s 模型导出为 ONNX、TorchScript 和 CoreML 格式。yolov5s.pt是最轻、最快的型号。其他选项是yolov5m.pt, yolov5l.ptand yolov5x.pt, 或者您拥有训练自定义数据集的检查点runs/exp0/weights/best.pt。有关所有可用模型的详细信息，请参阅我们的自述文件表。

python models/export.py --weights yolov5s.pt --img 640 --batch 1  # export at 640x640 with batch size 1

输出：

Namespace(batch_size=1, device='cpu', dynamic=False, half=False, img_size=[640, 640], include=['torchscript', 'onnx', 'coreml'], inplace=False, optimize=False, simplify=False, train=True, weights='./yolov5s.pt')
YOLOv5 🚀 v5.0-87-gf12cef8 torch 1.8.1+cu101 CPU

Fusing layers... 
Model Summary: 224 layers, 7266973 parameters, 0 gradients

PyTorch: starting from ./yolov5s.pt (14.8 MB)

TorchScript: starting export with torch 1.8.1+cu101...
TorchScript: export success, saved as ./yolov5s.torchscript.pt (29.4 MB)

ONNX: starting export with onnx 1.9.0...
ONNX: export success, saved as ./yolov5s.onnx (29.1 MB)

CoreML: starting export with coremltools 4.1...
CoreML: export success, saved as ./yolov5s.mlmodel (29.1 MB)

Export complete (10.40s). Visualize with https://github.com/lutzroeder/netron.

3 个导出的模型将与原始 PyTorch 模型一起保存：

推荐使用Netron Viewer来可视化导出的模型：