TFLite、ONNX、CoreML、TensorRT 导出#
📚 本指南介绍如何将训练好的 YOLOv5 🚀 模型从 PyTorch 导出为多种部署格式,包括 ONNX、TensorRT、CoreML 等。
开始之前#
克隆仓库,并在包含 PyTorch>=1.8 的 Python>=3.8.0 环境中安装 requirements.txt。模型和数据集会从 YOLOv5 最新版本自动下载。
git clone https://github.com/ultralytics/yolov5 # clone
cd yolov5
pip install -r requirements.txt # install有关 TensorRT 导出示例(需要 GPU),请参阅 Colab 笔记本的附录部分。
支持的导出格式#
YOLOv5 官方支持以下 12 种格式的推理:
| 格式 | export.py --include | 模型 |
|---|---|---|
| PyTorch | - | yolov5s.pt |
| TorchScript | torchscript | yolov5s.torchscript |
| ONNX | onnx | yolov5s.onnx |
| OpenVINO | openvino | yolov5s_openvino_model/ |
| TensorRT | engine | yolov5s.engine |
| CoreML | coreml | yolov5s.mlpackage |
| TensorFlow SavedModel | saved_model | yolov5s_saved_model/ |
| TensorFlow GraphDef | pb | yolov5s.pb |
| TensorFlow Lite | tflite | yolov5s.tflite |
| TensorFlow Edge TPU | edgetpu | yolov5s_edgetpu.tflite |
| TensorFlow.js | tfjs | yolov5s_web_model/ |
| PaddlePaddle | paddle | yolov5s_paddle_model/ |
基准测试#
以下基准测试在 Colab Pro 上使用 YOLOv5 教程笔记本 运行。复现步骤:
python benchmarks.py --weights yolov5s.pt --imgsz 640 --device 0Colab Pro V100 GPU#
benchmarks: weights=/content/yolov5/yolov5s.pt, imgsz=640, batch_size=1, data=/content/yolov5/data/coco128.yaml, device=0, half=False, test=False
Checking setup...
YOLOv5 🚀 v6.1-135-g7926afc torch 1.10.0+cu111 CUDA:0 (Tesla V100-SXM2-16GB, 16160MiB)
Setup complete ✅ (8 CPUs, 51.0 GB RAM, 46.7/166.8 GB disk)
Benchmarks complete (458.07s)
Format mAP@0.5:0.95 Inference time (ms)
0 PyTorch 0.4623 10.19
1 TorchScript 0.4623 6.85
2 ONNX 0.4623 14.63
3 OpenVINO NaN NaN
4 TensorRT 0.4617 1.89
5 CoreML NaN NaN
6 TensorFlow SavedModel 0.4623 21.28
7 TensorFlow GraphDef 0.4623 21.22
8 TensorFlow Lite NaN NaN
9 TensorFlow Edge TPU NaN NaN
10 TensorFlow.js NaN NaNColab Pro CPU#
benchmarks: weights=/content/yolov5/yolov5s.pt, imgsz=640, batch_size=1, data=/content/yolov5/data/coco128.yaml, device=cpu, half=False, test=False
Checking setup...
YOLOv5 🚀 v6.1-135-g7926afc torch 1.10.0+cu111 CPU
Setup complete ✅ (8 CPUs, 51.0 GB RAM, 41.5/166.8 GB disk)
Benchmarks complete (241.20s)
Format mAP@0.5:0.95 Inference time (ms)
0 PyTorch 0.4623 127.61
1 TorchScript 0.4623 131.23
2 ONNX 0.4623 69.34
3 OpenVINO 0.4623 66.52
4 TensorRT NaN NaN
5 CoreML NaN NaN
6 TensorFlow SavedModel 0.4623 123.79
7 TensorFlow GraphDef 0.4623 121.57
8 TensorFlow Lite 0.4623 316.61
9 TensorFlow Edge TPU NaN NaN
10 TensorFlow.js NaN NaN导出训练好的 YOLOv5 模型#
此命令会将预训练的 YOLOv5s 模型导出为 TorchScript 和 ONNX 格式。yolov5s.pt 是“small”模型,也是可用模型中第二小的一个。其他选项包括 yolov5n.pt、yolov5m.pt、yolov5l.pt 和 yolov5x.pt,以及对应的 P6 版本,例如 yolov5s6.pt,或你自己的自定义训练检查点,例如 runs/train/exp/weights/best.pt。有关所有可用模型的详情,请参阅 README 表格。
python export.py --weights yolov5s.pt --include torchscript onnx添加 --half,即可使用 FP16 半精度精度导出模型,以减小文件大小
输出:
export: data=data/coco128.yaml, weights=['yolov5s.pt'], imgsz=[640, 640], batch_size=1, device=cpu, half=False, inplace=False, train=False, keras=False, optimize=False, int8=False, dynamic=False, simplify=False, opset=12, verbose=False, workspace=4, nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, conf_thres=0.25, include=['torchscript', 'onnx']
YOLOv5 🚀 v6.2-104-ge3e5122 Python-3.8.0 torch-1.12.1+cu113 CPU
Downloading https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5s.pt to yolov5s.pt...
100% 14.1M/14.1M [00:00<00:00, 274MB/s]
Fusing layers...
YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients
PyTorch: starting from yolov5s.pt with output shape (1, 25200, 85) (14.1 MB)
TorchScript: starting export with torch 1.12.1+cu113...
TorchScript: export success ✅ 1.7s, saved as yolov5s.torchscript (28.1 MB)
ONNX: starting export with onnx 1.12.0...
ONNX: export success ✅ 2.3s, saved as yolov5s.onnx (28.0 MB)
Export complete (5.5s)
Results saved to /content/yolov5
Detect: python detect.py --weights yolov5s.onnx
Validate: python val.py --weights yolov5s.onnx
PyTorch Hub: model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.onnx')
Visualize: https://netron.app/导出的模型将与原始 PyTorch 模型保存在同一目录中:

建议使用 Netron Viewer 可视化导出的模型:

导出模型的使用示例#
detect.py 对导出的模型运行推理:
python detect.py --weights yolov5s.pt # PyTorch
python detect.py --weights yolov5s.torchscript # TorchScript
python detect.py --weights yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
python detect.py --weights yolov5s_openvino_model # OpenVINO
python detect.py --weights yolov5s.engine # TensorRT
python detect.py --weights yolov5s.mlpackage # CoreML (macOS only)
python detect.py --weights yolov5s_saved_model # TensorFlow SavedModel
python detect.py --weights yolov5s.pb # TensorFlow GraphDef
python detect.py --weights yolov5s.tflite # TensorFlow Lite
python detect.py --weights yolov5s_edgetpu.tflite # TensorFlow Edge TPU
python detect.py --weights yolov5s_paddle_model # PaddlePaddleval.py 对导出的模型运行验证:
python val.py --weights yolov5s.pt # PyTorch
python val.py --weights yolov5s.torchscript # TorchScript
python val.py --weights yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
python val.py --weights yolov5s_openvino_model # OpenVINO
python val.py --weights yolov5s.engine # TensorRT
python val.py --weights yolov5s.mlpackage # CoreML (macOS only)
python val.py --weights yolov5s_saved_model # TensorFlow SavedModel
python val.py --weights yolov5s.pb # TensorFlow GraphDef
python val.py --weights yolov5s.tflite # TensorFlow Lite
python val.py --weights yolov5s_edgetpu.tflite # TensorFlow Edge TPU
python val.py --weights yolov5s_paddle_model # PaddlePaddle将 PyTorch Hub 与导出的 YOLOv5 模型配合使用:
import torch
# Model
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.pt")
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.torchscript") # TorchScript
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.onnx") # ONNX Runtime
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s_openvino_model") # OpenVINO
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.engine") # TensorRT
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.mlpackage") # CoreML(仅限 macOS)
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s_saved_model") # TensorFlow SavedModel
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.pb") # TensorFlow GraphDef
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.tflite") # TensorFlow Lite
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s_edgetpu.tflite") # TensorFlow Edge TPU
model = torch.hub.load("ultralytics/yolov5", "custom", "yolov5s_paddle_model") # PaddlePaddle
# Images
img = "https://ultralytics.com/images/zidane.jpg" # 或文件、Path、PIL、OpenCV、numpy、列表
# Inference
results = model(img)
# Results
results.print() # 或 .show(), .save(), .crop(), .pandas() 等。OpenCV DNN 推理#
使用 ONNX 模型进行 OpenCV 推理:
python export.py --weights yolov5s.pt --include onnx
python detect.py --weights yolov5s.onnx --dnn # detect
python val.py --weights yolov5s.onnx --dnn # validateC++ 推理#
YOLOv5 OpenCV DNN C++ 推理导出的 ONNX 模型示例:
- https://github.com/Hexmagic/ONNX-yolov5/blob/master/src/test.cpp
- https://github.com/doleron/yolov5-opencv-cpp-python
YOLOv5 OpenVINO C++ 推理示例:
- https://github.com/dacquaviva/yolov5-openvino-cpp-python
- https://github.com/UNeedCryDear/yolov5-seg-opencv-onnxruntime-cpp
TensorFlow.js 网页浏览器推理#
支持的环境#
Ultralytics 提供多种开箱即用的环境,每种环境都预装了 CUDA、CUDNN、Python 和 PyTorch 等必备依赖,助你快速启动项目。
- 免费 GPU 笔记本:
- Google Cloud:GCP 快速入门指南
- Amazon:AWS 快速入门指南
- Azure:AzureML 快速入门指南
- Docker:Docker 快速入门指南
项目状态#
此徽章表示所有 YOLOv5 GitHub Actions 持续集成 (CI) 测试均已成功通过。这些 CI 测试会严格检查 YOLOv5 在多个关键方面的功能和性能:训练、验证、推理、导出和基准测试。这些测试每 24 小时以及每次提交新代码后运行一次,确保 YOLOv5 在 macOS、Windows 和 Ubuntu 上保持稳定可靠地运行。
常见问题#
CPU 部署使用 ONNX 或 OpenVINO,其运行速度最高可比 PyTorch 快 3 倍;NVIDIA GPU 部署使用 TensorRT,其运行速度最高可快 5 倍。CoreML 面向 Apple 设备,TensorFlow Lite 或 Edge TPU 则面向 Android 和 Coral 硬件。
添加
--half。大多数格式的半精度导出都需要 GPU,因此还要传入--device 0;CoreML 是个例外,可以在 CPU 上导出 FP16。可以。添加
--dynamic。ONNX 导出会启用动态批次、高度和宽度轴,并在 CPU 上进行跟踪,因此请省略--device 0。TensorRT 导出会保持图像大小固定,仅改变批次维度:仍然需要--device 0,而--batch-size用于设置引擎可接受的最大批次大小。将导出的文件传递给
detect.py --weights、val.py --weights或torch.hub.load("ultralytics/yolov5", "custom", "yolov5s.onnx")。YOLOv5 会根据文件后缀选择正确的后端。