Exportation TFLite, ONNX, CoreML, TensorRT
📚 Ce guide explique comment exporter un modèle YOLOv5 🚀 entraîné depuis PyTorch vers différents formats de déploiement, notamment ONNX, TensorRT, CoreML et autres.
Avant de commencer
Clonez le dépôt et installez requirements.txt dans un environnement Python>=3.8.0, incluant PyTorch>=1.8. Les modèles et les ensembles de données sont téléchargés automatiquement à partir de la dernière version de YOLOv5.
git clone https://github.com/ultralytics/yolov5 # clone
cd yolov5
pip install -r requirements.txt # install
Pour TensorRT exemple d'exportation (nécessite un GPU), voir notre Colab notebook section annexe.
Formats d'exportation pris en charge
L'inférence YOLOv5 est officiellement prise en charge dans 12 formats :
Conseils relatifs aux performances
- Exporter vers ONNX ou OpenVINO pour un gain de vitesse CPU jusqu'à 3x. Voir les Bancs d'essai CPU.
- Exporter vers TensorRT pour un gain de vitesse GPU jusqu'à 5x. Voir les Bancs d'essai GPU.
Format | export.py --include |
Modèle |
---|---|---|
PyTorch | - | yolov5s.pt |
TorchScript | torchscript |
yolov5s.torchscript |
ONNX | onnx |
yolov5s.onnx |
OpenVINO | openvino |
yolov5s_openvino_model/ |
TensorRT | engine |
yolov5s.engine |
CoreML | coreml |
yolov5s.mlmodel |
SavedModel TensorFlow | saved_model |
yolov5s_saved_model/ |
GraphDef TensorFlow | 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/ |
Bancs d'essai
Les benchmarks ci-dessous sont exécutés sur un Colab Pro avec le notebook tutoriel YOLOv5 . Pour reproduire:
python benchmarks.py --weights yolov5s.pt --imgsz 640 --device 0
GPU Colab Pro V100
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 NaN
CPU Colab Pro
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
Exporter un modèle YOLOv5 entraîné
Cette commande exporte un modèle YOLOv5s pré-entraîné aux formats TorchScript et ONNX. yolov5s.pt
est le modèle 'small', le deuxième plus petit modèle disponible. Les autres options sont yolov5n.pt
, yolov5m.pt
, yolov5l.pt
et yolov5x.pt
, ainsi que leurs équivalents P6, c'est-à-dire yolov5s6.pt
ou votre propre point de contrôle d'entraînement personnalisé, c'est-à-dire runs/exp/weights/best.pt
. Veuillez consulter notre fichier README pour obtenir des informations détaillées sur tous les modèles disponibles. table.
python export.py --weights yolov5s.pt --include torchscript onnx
Astuce
Ajouter --half
pour exporter les modèles en FP16 half précision pour des fichiers de plus petite taille
Sortie :
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/
Les 3 modèles exportés seront enregistrés à côté du modèle PyTorch original :
Netron Viewer est recommandé pour visualiser les modèles exportés :
Exemples d'utilisation du modèle exporté
detect.py
exécute l'inférence sur les modèles exportés :
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=True
python detect.py --weights yolov5s_openvino_model # OpenVINO
python detect.py --weights yolov5s.engine # TensorRT
python detect.py --weights yolov5s.mlmodel # 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 # PaddlePaddle
val.py
exécute la validation sur les modèles exportés :
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=True
python val.py --weights yolov5s_openvino_model # OpenVINO
python val.py --weights yolov5s.engine # TensorRT
python val.py --weights yolov5s.mlmodel # 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
Utilisez PyTorch Hub avec les modèles YOLOv5 exportés :
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.mlmodel") # CoreML (macOS Only)
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" # or file, Path, PIL, OpenCV, numpy, list
# Inference
results = model(img)
# Results
results.print() # or .show(), .save(), .crop(), .pandas(), etc.
Inférence OpenCV DNN
Inférence OpenCV avec les modèles ONNX :
python export.py --weights yolov5s.pt --include onnx
python detect.py --weights yolov5s.onnx --dnn # detect
python val.py --weights yolov5s.onnx --dnn # validate
Inférence C++
Exemples d'inférence YOLOv5 OpenCV DNN C++ sur des modèles ONNX exportés :
- https://github.com/Hexmagic/ONNX-yolov5/blob/master/src/test.cpp
- https://github.com/doleron/yolov5-opencv-cpp-python
Exemples d'inférence YOLOv5 OpenVINO C++ :
- https://github.com/dacquaviva/yolov5-openvino-cpp-python
- https://github.com/UNeedCryDear/yolov5-seg-opencv-dnn-cpp
Inférence avec navigateur web TensorFlow.js
Environnements pris en charge
Ultralytics propose une gamme d'environnements prêts à l'emploi, chacun préinstallé avec les dépendances essentielles telles que CUDA, CUDNN, Python et PyTorch, pour démarrer vos projets.
- Notebooks GPU Gratuits:
- Google Cloud : Guide de démarrage rapide GCP
- Amazon : Guide de démarrage rapide AWS
- Azure : Guide de démarrage rapide AzureML
- Docker: Guide de démarrage rapide Docker
État du projet
Ce badge indique que tous les tests d'intégration continue (CI) YOLOv5 GitHub Actions réussissent. Ces tests CI vérifient rigoureusement la fonctionnalité et les performances de YOLOv5 à travers divers aspects clés : l'entraînement, la validation, l'inférence, l'exportation et les benchmarks. Ils garantissent un fonctionnement cohérent et fiable sur macOS, Windows et Ubuntu, avec des tests effectués toutes les 24 heures et à chaque nouveau commit.