Vai al contenuto

TFLite, ONNX, CoreML, TensorRT Esportazione

📚 Questa guida spiega come esportare un modello addestrato YOLOv5 🚀 da PyTorch ai formati ONNX e TorchScript .

Prima di iniziare

Clona il repo e installa il file requirements.txt in un file di tipo Python>=3.8.0 con l'inclusione di PyTorch>=1.8. I modelli e i dataset vengono scaricati automaticamente dall'ultimarelease di YOLOv5 .

git clone https://github.com/ultralytics/yolov5  # clone
cd yolov5
pip install -r requirements.txt  # install

Per TensorRT esempio di esportazione (richiede una GPU) vedi il nostro Colab taccuino sezione appendice. Apri in Colab

Formati

YOLOv5 è ufficialmente supportata in 11 formati:

💡 Suggerimento: esporta in ONNX o OpenVINO per aumentare la velocità della CPU fino a 3 volte. Vedi Benchmark della CPU. 💡 Suggerimento: esporta in TensorRT per ottenere una velocità della GPU fino a 5 volte superiore. Vedi Benchmark GPU.

Formato export.py --include Modello
PyTorch - yolov5s.pt
TorchScript torchscript yolov5s.torchscript
ONNX onnx yolov5s.onnx
OpenVINO openvino yolov5s_openvino_model/
TensorRT engine yolov5s.engine
CoreML coreml yolov5s.mlmodel
TensorFlow SavedModel saved_model yolov5s_saved_model/
TensorFlow GraphDef pb yolov5s.pb
TensorFlow Lite tflite yolov5s.tflite
TensorFlow Bordo TPU edgetpu yolov5s_edgetpu.tflite
TensorFlow.js tfjs yolov5s_web_model/
PaddlePaddle paddle yolov5s_paddle_model/

Parametri di riferimento

I benchmark riportati di seguito sono stati eseguiti su un Colab Pro con il notebook tutorial YOLOv5 Apri in Colab. Per riprodurre:

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

Esporta un modello addestrato di YOLOv5

Questo comando esporta un modello YOLOv5s pre-addestrato nei formati TorchScript e ONNX . yolov5s.pt è il modello "piccolo", il secondo più piccolo disponibile. Le altre opzioni sono yolov5n.pt, yolov5m.pt, yolov5l.pt e yolov5x.pt, insieme alle loro controparti P6, ovvero yolov5s6.pt o il tuo checkpoint di formazione personalizzato, ad esempio. runs/exp/weights/best.pt. Per maggiori dettagli su tutti i modelli disponibili, consulta il nostro README. Tavolo.

python export.py --weights yolov5s.pt --include torchscript onnx

💡 Suggerimento: Aggiungi --half esportare i modelli a mezza precisione FP16 per ridurre le dimensioni dei file

Uscita:

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/

I 3 modelli esportati saranno salvati insieme al modello originale PyTorch :

YOLO luoghi di esportazione

Netron Viewer è consigliato per visualizzare i modelli esportati:

YOLO visualizzazione del modello

Esempi di utilizzo del modello esportato

detect.py esegue l'inferenza sui modelli esportati:

python detect.py --weights yolov5s.pt                 # PyTorch
                           yolov5s.torchscript        # TorchScript
                           yolov5s.onnx               # ONNX Runtime or OpenCV DNN with dnn=True
                           yolov5s_openvino_model     # OpenVINO
                           yolov5s.engine             # TensorRT
                           yolov5s.mlmodel            # CoreML (macOS only)
                           yolov5s_saved_model        # TensorFlow SavedModel
                           yolov5s.pb                 # TensorFlow GraphDef
                           yolov5s.tflite             # TensorFlow Lite
                           yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
                           yolov5s_paddle_model       # PaddlePaddle

val.py esegue la convalida dei modelli esportati:

python val.py --weights yolov5s.pt                 # PyTorch
                        yolov5s.torchscript        # TorchScript
                        yolov5s.onnx               # ONNX Runtime or OpenCV DNN with dnn=True
                        yolov5s_openvino_model     # OpenVINO
                        yolov5s.engine             # TensorRT
                        yolov5s.mlmodel            # CoreML (macOS Only)
                        yolov5s_saved_model        # TensorFlow SavedModel
                        yolov5s.pb                 # TensorFlow GraphDef
                        yolov5s.tflite             # TensorFlow Lite
                        yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
                        yolov5s_paddle_model       # PaddlePaddle

Usa PyTorch Hub con i modelli esportati di 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.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.

Inferenza DNN di OpenCV

Inferenza in OpenCV con i modelli di 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

Inferenza in C++

YOLOv5 Inferenza OpenCV DNN C++ su esempi di modelli esportati da ONNX :

YOLOv5 OpenVINO Esempi di inferenza in C++:

TensorFlowInferenza del browser web .js

Ambienti supportati

Ultralytics offre una serie di ambienti pronti all'uso, ciascuno preinstallato con le dipendenze essenziali come CUDA, CUDNN, Python, e PyTorchper dare il via ai tuoi progetti.

Stato del progetto

YOLOv5 CI

Questo badge indica che tutti i test di YOLOv5 GitHub Actions Continuous Integration (CI) sono stati superati con successo. Questi test CI verificano rigorosamente la funzionalità e le prestazioni di YOLOv5 in vari aspetti chiave: formazione, validazione, inferenza, esportazione e benchmark. Assicurano un funzionamento coerente e affidabile su macOS, Windows e Ubuntu, con test condotti ogni 24 ore e su ogni nuovo commit.



Creato 2023-11-12, Aggiornato 2024-01-12
Autori: glenn-jocher (5)

Commenti