Exportação para TFLite, ONNX, CoreML e TensorRT#
📚 Este guia explica como exportar um modelo YOLOv5 🚀 treinado de PyTorch para vários formatos de implementação, incluindo ONNX, TensorRT, CoreML e outros.
Antes de começar#
Clone o repositório e instale requirements.txt em um ambiente Python>=3.8.0, incluindo PyTorch>=1.8. Os modelos e conjuntos de dados são baixados automaticamente da versão mais recente do YOLOv5.
git clone https://github.com/ultralytics/yolov5 # clone
cd yolov5
pip install -r requirements.txt # installPara ver um exemplo de exportação para TensorRT (requer GPU), consulta a secção do apêndice do notebook Colab.
Formatos de exportação compatíveis#
A inferência YOLOv5 é oficialmente compatível com 12 formatos:
- Exporta para ONNX ou OpenVINO para obter até 3 vezes mais velocidade na CPU. Consulta os benchmarks de CPU.
- Exporta para TensorRT para obter até 5 vezes mais velocidade na GPU. Consulta os benchmarks de 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.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/ |
Benchmarks#
Os benchmarks abaixo são executados no Colab Pro com o notebook do tutorial YOLOv5 . Para reproduzi-los:
python benchmarks.py --weights yolov5s.pt --imgsz 640 --device 0GPU V100 do 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 NaNCPU do 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 NaNExportar 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 disponível. As outras opções são yolov5n.pt, yolov5m.pt, yolov5l.pt e yolov5x.pt, bem como as respetivas versões P6, ou seja, yolov5s6.pt, ou o teu próprio checkpoint de treino personalizado, ou seja, runs/train/exp/weights/best.pt. Para ver informações sobre todos os modelos disponíveis, consulta a tabela do nosso README.
python export.py --weights yolov5s.pt --include torchscript onnxAdiciona --half para exportar modelos com precisão FP16 de meia precisão, reduzindo o tamanho dos ficheiros.
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 modelos exportados serão guardados junto ao modelo PyTorch original:

Recomendamos o Netron Viewer para visualizar os modelos exportados:

Exemplos de utilização de modelos exportados#
detect.py executa inferências 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
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 executa a validação de 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
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 # PaddlePaddleUsa 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.mlpackage") # CoreML (apenas 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" # ou file, Path, PIL, OpenCV, numpy, list
# Inference
results = model(img)
# Results
results.print() # ou .show(), .save(), .crop(), .pandas(), etc.Inferência com OpenCV DNN#
Inferência com modelos ONNX usando OpenCV:
python export.py --weights yolov5s.pt --include onnx
python detect.py --weights yolov5s.onnx --dnn # detect
python val.py --weights yolov5s.onnx --dnn # validateInferência em C++#
Exemplos de inferência em C++ com OpenCV DNN do YOLOv5 num modelo ONNX exportado:
- https://github.com/Hexmagic/ONNX-yolov5/blob/master/src/test.cpp
- https://github.com/doleron/yolov5-opencv-cpp-python
Exemplos de inferência em C++ com OpenVINO do YOLOv5:
- https://github.com/dacquaviva/yolov5-openvino-cpp-python
- https://github.com/UNeedCryDear/yolov5-seg-opencv-onnxruntime-cpp
Inferência no navegador com TensorFlow.js#
Ambientes compatíveis#
A Ultralytics oferece uma variedade de ambientes prontos para uso, cada um com dependências essenciais pré-instaladas, como CUDA, CUDNN, Python e PyTorch, para você começar seus projetos.
- Notebooks gratuitos com GPU:
- 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
Estado do projeto#
Este emblema indica que todos os testes de Integração contínua (CI) das GitHub Actions da YOLOv5 estão a passar. Estes testes de CI verificam rigorosamente a funcionalidade e o desempenho da YOLOv5 em vários aspetos importantes: treino, validação, inferência, exportação e avaliação de desempenho. Os testes garantem um funcionamento consistente e fiável em macOS, Windows e Ubuntu e são realizados a cada 24 horas e sempre que é criado um novo commit.
Perguntas frequentes#
Usa ONNX ou OpenVINO para implementação em CPU, onde são até 3 vezes mais rápidos do que o PyTorch, e TensorRT para GPUs NVIDIA, onde é até 5 vezes mais rápido. O CoreML destina-se a dispositivos Apple, e o TensorFlow Lite ou o Edge TPU destinam-se a dispositivos Android e hardware Coral.
Adiciona
--half. A meia precisão requer uma GPU para a maioria dos formatos, por isso passa também--device 0; o CoreML é a exceção e pode exportar em FP16 na CPU.Sim. Adiciona
--dynamic. As exportações ONNX passam a ter eixos dinâmicos para o lote, a altura e a largura e são rastreadas na CPU, por isso não incluas--device 0. As exportações TensorRT mantêm fixo o tamanho da imagem e variam apenas a dimensão do lote: continuam a precisar de--device 0, e--batch-sizedefine o maior lote que o motor aceita.Passa o ficheiro exportado a
detect.py --weights,val.py --weightsoutorch.hub.load("ultralytics/yolov5", "custom", "yolov5s.onnx"). O YOLOv5 seleciona o backend correto com base na extensão do ficheiro.