TFLite, ONNX, CoreML, TensorRT Exportar

Este guia explica como exportar um modelo treinado YOLOv5 🚀 de PyTorch para os formatos ONNX e TorchScript .

Antes de começares

Clona o repo e instala o requirements.txt em um Python>=3.8.0 incluindo PyTorch>=1.8. Os modelos e conjuntos de dados são descarregados automaticamente a partir daversão mais recente de YOLOv5 .

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

Para TensorRT exemplo de exportação (requer GPU) vê o nosso Colab caderno secção do anexo.

Formatos

YOLOv5 A inferência é oficialmente suportada em 11 formatos:

Dica: exporta para ONNX ou OpenVINO para aumentar a velocidade da CPU até 3x. Vê os Indicadores de desempenho da CPU. Sugestão: exporta para TensorRT para obteres uma aceleração da GPU até 5x. Vê os Indicadores de desempenho da GPU.

Formato	`export.py --include`	Modelo
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 Leve	`tflite`	`yolov5s.tflite`
TensorFlow Borda TPU	`edgetpu`	`yolov5s_edgetpu.tflite`
TensorFlow.js	`tfjs`	`yolov5s_web_model/`
PaddlePaddle	`paddle`	`yolov5s_paddle_model/`

Referências

Os testes de referência abaixo foram executados num Colab Pro com o notebook tutorial YOLOv5 . Para reproduzir:

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

Colab 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

Exporta um modelo YOLOv5 treinado

Este comando exporta um modelo YOLOv5s pré-treinado para os formatos TorchScript e ONNX . yolov5s.pt é o modelo "pequeno", o segundo modelo mais pequeno disponível. As outras opções são yolov5n.pt, yolov5m.pt, yolov5l.pt e yolov5x.ptjuntamente com os seus homólogos P6, ou seja yolov5s6.pt ou o teu próprio ponto de controlo de formação personalizado, ou seja runs/exp/weights/best.pt. Para mais informações sobre todos os modelos disponíveis, consulta o nosso README tabela.

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

💡 ProTip: Adiciona --half para exportar modelos com meia precisão FP16 para ficheiros mais pequenos

Saída:

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/

Os 3 modelos exportados serão guardados juntamente com o modelo original PyTorch :

YOLO locais de exportação

O Netron Viewer é recomendado para visualizar os modelos exportados:

YOLO visualização do modelo

Exemplos de utilização de modelos exportados

detect.py executa a inferência nos modelos exportados:

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 executa a validação dos modelos exportados:

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

Utiliza o PyTorch Hub com modelos YOLOv5 exportados:

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.

Inferência OpenCV DNN

Inferência OpenCV com modelos 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

Inferência C++

YOLOv5 Inferência OpenCV DNN C++ em exemplos de modelos ONNX exportados:

YOLOv5 OpenVINO Exemplos de inferência C++:

TensorFlowInferência de navegador da Web .js

https://aukerul-shuvo.github.io/YOLOv5_TensorFlow-JS/

Ambientes suportados

Ultralytics fornece uma gama de ambientes prontos a utilizar, cada um pré-instalado com dependências essenciais como CUDA, CUDNN, Python, e PyTorchpara dar início aos teus projectos.

Portáteis GPU gratuitos:
Google Cloud: Guia de início rápido do GCP
Amazon: Guia de início rápido do AWS
Azure: Guia de início rápido do AzureML
Docker: Guia de início rápido do Docker

Estado do projeto

Este emblema indica que todos os testes de Integração Contínua (CI) do YOLOv5 GitHub Actions foram aprovados com êxito. Esses testes de CI verificam rigorosamente a funcionalidade e o desempenho do YOLOv5 em vários aspectos principais: treinamento, validação, inferência, exportação e benchmarks. Eles garantem uma operação consistente e confiável no macOS, Windows e Ubuntu, com testes realizados a cada 24 horas e a cada novo commit.

Criado em 2023-11-12, Atualizado em 2024-01-12
Autores: glenn-jocher (5)