Exportação TFLite, ONNX, CoreML, TensorRT#

📚 Este guia explica como exportar um modelo YOLOv5 🚀 treinado do PyTorch para vários formatos de implantação, incluindo ONNX, TensorRT, CoreML e mais.

Antes de começar#

Clona o repositório e instala o requirements.txt num ambiente Python>=3.8.0, incluindo o PyTorch>=1.8. Os Models e os datasets são descarregados automaticamente a partir do mais recente release do YOLOv5.

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

Para um exemplo de exportação para TensorRT (requer GPU), consulta a seção de apêndice do nosso notebook do Colab.

Formatos de Exportação Suportados#

A inferência do YOLOv5 é oficialmente suportada em 12 formatos:

Dicas de desempenho

Exporta para ONNX ou OpenVINO para obter até 3x mais velocidade em CPU. Consulta os CPU Benchmarks.
Exporta para TensorRT para obter até 5x mais velocidade em GPU. Consulta os GPU Benchmarks.

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 Lite	`tflite`	`yolov5s.tflite`
TensorFlow Edge TPU	`edgetpu`	`yolov5s_edgetpu.tflite`
TensorFlow.js	`tfjs`	`yolov5s_web_model/`
PaddlePaddle	`paddle`	`yolov5s_paddle_model/`

Benchmarks#

Os benchmarks abaixo rodam em um Colab Pro com o notebook de tutorial do YOLOv5 . Para reproduzir:

python benchmarks.py --weights yolov5s.pt --imgsz 640 --device 0

GPU V100 no Colab Pro#

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 no 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

Exportar um modelo YOLOv5 treinado#

Este comando exporta um modelo YOLOv5s pré-treinado para os formatos TorchScript e ONNX. yolov5s.pt é o modelo 'small', o segundo menor modelo disponível. Outras opções são yolov5n.pt, yolov5m.pt, yolov5l.pt e yolov5x.pt, juntamente com as suas contrapartes P6, ou seja, yolov5s6.pt, ou o teu próprio ponto de controle de treino personalizado, ou seja, runs/exp/weights/best.pt. Para detalhes sobre todos os modelos disponíveis, consulta a tabela do nosso README.

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

Dica

Adiciona --half para exportar modelos com precisão reduzida FP16 half para tamanhos de ficheiro menores

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 salvos junto ao modelo PyTorch original:

YOLO export locations

O Netron Viewer é recomendado para visualizar modelos exportados:

YOLO model visualization

Exemplos de uso de modelo exportado#

detect.py executa inferência em modelos exportados:

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

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

Use 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 com OpenCV utilizando 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 em C++#

Exemplos de inferência YOLOv5 OpenCV DNN C++ em modelo ONNX exportado:

Exemplos de inferência YOLOv5 OpenVINO C++:

Inferência no navegador web TensorFlow.js#

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

Ambientes suportados#

A Ultralytics disponibiliza uma gama de ambientes prontos a usar, cada um pré-instalado com dependências essenciais como CUDA, CUDNN, Python e PyTorch, para dar início aos teus projetos.

Notebooks de GPU Grátis:
Google Cloud: Guia de Início Rápido do GCP
Amazon: Guia de Início Rápido da AWS
Azure: Guia de Início Rápido do AzureML
Docker: Guia de Início Rápido do Docker

Status do projeto#

Este crachá indica que todos os testes de Integração Contínua (CI) do GitHub Actions do YOLOv5 estão a passar com sucesso. Estes testes de CI verificam rigorosamente a funcionalidade e o desempenho do YOLOv5 em vários aspetos fundamentais: treino, validação, inferência, exportação e testes de desempenho. Garantem um funcionamento consistente e fiável no macOS, Windows e Ubuntu, com testes realizados a cada 24 horas e a cada novo commit.

Contribuidores

GLglenn-jocher²

Criado há 2 mesesAtualizado mês passado