TFLite, ONNX, CoreML, TensorRT Export

[glenn-jocher]

📚 This guide explains how to export a trained YOLOv5 🚀 model from PyTorch to ONNX and TorchScript formats. UPDATED 8 December 2022.

Before You Start

Clone repo and install requirements.txt in a Python>=3.7.0 environment, including PyTorch>=1.7. Models and datasets download automatically from the latest YOLOv5 release.

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

For TensorRT export example (requires GPU) see our Colab notebook appendix section.

Formats

YOLOv5 inference is officially supported in 11 formats:

💡 ProTip: Export to ONNX or OpenVINO for up to 3x CPU speedup. See CPU Benchmarks.
💡 ProTip: Export to TensorRT for up to 5x GPU speedup. See GPU Benchmarks.

Format	export.py --include	Model
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

Benchmarks below run on a Colab Pro with the YOLOv5 tutorial notebook . To reproduce:

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

Colab Pro V100 GPU

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

Export a Trained YOLOv5 Model

This command exports a pretrained YOLOv5s model to TorchScript and ONNX formats. yolov5s.pt is the 'small' model, the second smallest model available. Other options are yolov5n.pt, yolov5m.pt, yolov5l.pt and yolov5x.pt, along with their P6 counterparts i.e. yolov5s6.pt or you own custom training checkpoint i.e. runs/exp/weights/best.pt. For details on all available models please see our README table.

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

💡 ProTip: Add --half to export models at FP16 half precision for smaller file sizes

Output:

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.7.13 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/

The 3 exported models will be saved alongside the original PyTorch model:

Netron Viewer is recommended for visualizing exported models:

Exported Model Usage Examples

detect.py runs inference on exported models:

python detect.py --weights yolov5s.pt                 # PyTorch
                           yolov5s.torchscript        # TorchScript
                           yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                           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 runs validation on exported models:

python val.py --weights yolov5s.pt                 # PyTorch
                        yolov5s.torchscript        # TorchScript
                        yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                        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

Use PyTorch Hub with exported YOLOv5 models:

import torch

# Model
model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt')
                                                       'yolov5s.torchscript ')       # TorchScript
                                                       'yolov5s.onnx')               # ONNX Runtime
                                                       '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

# 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.

OpenCV DNN inference

OpenCV inference with ONNX models:

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

python detect.py --weights yolov5s.onnx --dnn  # detect
python val.py --weights yolov5s.onnx --dnn  # validate

C++ Inference

YOLOv5 OpenCV DNN C++ inference on exported ONNX model examples:

• https://github.com/Hexmagic/ONNX-yolov5/blob/master/src/test.cpp
• https://github.com/doleron/yolov5-opencv-cpp-python

YOLOv5 OpenVINO C++ inference examples:

• https://github.com/dacquaviva/yolov5-openvino-cpp-python
• https://github.com/UNeedCryDear/yolov5-seg-opencv-dnn-cpp

TensorFlow.js Web Browser Inference

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

Environments

YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

Status

If this badge is green, all YOLOv5 GitHub Actions Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv5 training, validation, inference, export and benchmarks on macOS, Windows, and Ubuntu every 24 hours and on every commit.

[TommyZihao]

Thank you so much!
I will deploy onnx model on mobile devices!

[tienhoang1094]

it only work with 5s pretrained,

[rcg12387]

@glenn-jocher My onnx is 1.7.0, python is 3.8.3, pytorch is 1.4.0 (your latest recommendation is 1.5.0).
But exporting to ONNX is failed because of opset version 12. This is my command line:

export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov5s.pt --img 640 --batch 1

And it failed with this error:

Fusing layers...
Model Summary: 140 layers, 7.45958e+06 parameters, 7.45958e+06 gradientsONNX export failed: Unsupported ONNX opset version: 12

I don't think it caused by PyTorch version lower than your recommendation.
Any advice? Thank you.

[rcg12387]

I changed opset_version to 11 in export.py, and new error messages came up:

Fusing layers...
Model Summary: 140 layers, 7.45958e+06 parameters, 7.45958e+06 gradients
Segmentation fault (core dumped)

This is the full message:

$ export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov5s.pt --img 640 --batch 1
Namespace(batch_size=1, img_size=[640, 640], weights='./weights/yolov5s.pt')
/home/DL-001/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/serialization.py:593: SourceChangeWarning: source code of class 'torch.nn.modules.conv.Conv2d' has changed. you can retrieve the original source code by accessing the object's source attribute or set `torch.nn.Module.dump_patches = True` and use the patch tool to revert the changes.
  warnings.warn(msg, SourceChangeWarning)
/home/DL-001/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/serialization.py:593: SourceChangeWarning: source code of class 'torch.nn.modules.container.ModuleList' has changed. you can retrieve the original source code by accessing the object's source attribute or set `torch.nn.Module.dump_patches = True` and use the patch tool to revert the changes.
  warnings.warn(msg, SourceChangeWarning)
TorchScript export failed: Only tensors or tuples of tensors can be output from traced functions (getOutput at /opt/conda/conda-bld/pytorch_1579022027550/work/torch/csrc/jit/tracer.cpp:212)
frame #0: c10::Error::Error(c10::SourceLocation, std::string const&) + 0x47 (0x7fb3a6bdf627 in /home/DL-001/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/lib/libc10.so)
frame #1: torch::jit::tracer::TracingState::getOutput(c10::IValue const&, unsigned long) + 0x334 (0x7fb3b16d2024 in /home/DL-001/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/lib/libtorch.so)
frame #2: torch::jit::tracer::trace(std::vector<c10::IValue, std::allocator<c10::IValue> >, std::function<std::vector<c10::IValue, std::allocator<c10::IValue> > (std::vector<c10::IValue, std::allocator<c10::IValue> >)> const&, std::function<std::string (at::Tensor const&)>, bool, torch::jit::script::Module*) + 0x539 (0x7fb3b16d99f9 in /home/DL-001/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/lib/libtorch.so)
frame #3: <unknown function> + 0x759fed (0x7fb3ddbcafed in /home/DL-001/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/lib/libtorch_python.so)
frame #4: <unknown function> + 0x7720ee (0x7fb3ddbe30ee in /home/DL-001/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/lib/libtorch_python.so)
frame #5: <unknown function> + 0x28b8a7 (0x7fb3dd6fc8a7 in /home/DL-001/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/lib/libtorch_python.so)
<omitting python frames>
frame #24: __libc_start_main + 0xe7 (0x7fb416e13b97 in /lib/x86_64-linux-gnu/libc.so.6)

Fusing layers...
Model Summary: 140 layers, 7.45958e+06 parameters, 7.45958e+06 gradients
Segmentation fault (core dumped)

[rcg12387]

I debugged it and found the reason.
It failed at ts = torch.jit.trace(model, img), so I realized it was caused by lower version of PyTorch.
Then I upgraded PyTorch to 1.5.1, and it worked good finally.

[Ezra-Yu]

why you set Detect() layer export=True? this will let Detect() layer not in the onnx model.

[glenn-jocher]

@Ezra-Yu yes that is correct. You are free to set it to False if that suits you better.

[ycdhqzhiai]

@glenn-jocher Why is the input of onnx fixed，but pt is multiple of 32

[neverrop]

hi, is there any sample code to use the exported onnx to get the Nx5 bbox?. i tried to use the postprocess from detect.py, but it doesnt work well.

[dlawrences]

hi, is there any sample code to use the exported onnx to get the Nx5 bbox?. i tried to use the postprocess from detect.py, but it doesnt work well.

Hi @neverrop

I have added guidance over how this could be achieved here: #343 (comment)

Hope this is useful!

[neverrop]

hi, is there any sample code to use the exported onnx to get the Nx5 bbox?. i tried to use the postprocess from detect.py, but it doesnt work well.

Hi @neverrop

I have added guidance over how this could be achieved here: #343 (comment)

Hope this is useful!.
Thank you so much. I will try it today｡

[shenglih]

Would CoreML failure as shown below affect the successfully converted onnx model? Thank you.

ONNX export success, saved as weights/yolov5s.onnx
WARNING:root:TensorFlow version 2.2.0 detected. Last version known to be fully compatible is 1.14.0 .
WARNING:root:Keras version 2.4.3 detected. Last version known to be fully compatible of Keras is 2.2.4 .

Starting CoreML export with coremltools 3.4...
CoreML export failure: module 'coremltools' has no attribute 'convert'

Export complete. Visualize with https://github.com/lutzroeder/netron

[dlawrences]

Hi @shenglih

CoreML export doesn't affect the ONNX one in any way.

Regards

[Mayur2992]

Starting CoreML export with coremltools 3.4...
CoreML export failure: module 'coremltools' has no attribute 'convert'

Export complete. Visualize with https://github.com/lutzroeder/netron.

anyone solved it?

[dlawrences]

Starting CoreML export with coremltools 3.4...
CoreML export failure: module 'coremltools' has no attribute 'convert'

Export complete. Visualize with https://github.com/lutzroeder/netron.

anyone solved it?

Hi. I think you need to update to the latest coremltools package version.

Please see this one: #315 (comment)

[zyyang]

Starting CoreML export with coremltools 3.4...
CoreML export failure: module 'coremltools' has no attribute 'convert'

Export complete. Visualize with https://github.com/lutzroeder/netron.

anyone solved it?

reinstall your coremltools:
pip install coremltools==4.0b2

[zhepherd]

pip install coremltools==4.0b2

my pytorch version is 1.4, coremltools=4.0b2,but error

Starting ONNX export with onnx 1.7.0...
Fusing layers... Model Summary: 284 layers, 8.84108e+07 parameters, 8.45317e+07 gradients
ONNX export failure: Unsupported ONNX opset version: 12

Starting CoreML export with coremltools 4.0b2...
CoreML export failure: name 'ts' is not defined
how to solved it

[dlawrences]

@zhepherd

Please install torch=1.5.1.

[Qausain]

@glenn-jocher Many thanks for your reply. I already gone through the link you sent. Actually my problem statement is that "I just want to create a very simple webpage where user can upload image & that image will store in s3 then my trained yolo model started work on that image & return result on the same webpage"? Your valuable answer in this regard will be very helpful for me Or you can refer me to any blogpost or video link. Thanks in advance

[zldrobit]

@Qausain Would you like trying this PR with Gradio integration. You could create a webpage with yolov5 prediction using a command line. However, this PR does not store the images.

[fzemi]

@glenn-jocher Any plans for adding ClassificationModel export to tflite?

[Qausain]

@zldrobit Thanks for your suggestion, Currently they only support models trained using the Ultralytics HUB. In mid January they may allow models trained locally to be imported into the Ultralytics HUB.

[mekidhamza]

I trained the yolov5 model on the custom dataset, but when I try to predict the bounding box, confidence, and the class name using python detect.py the result was perfect ( bounding box and confidence) but when I load the model with model = model = torch.hub.load('C:\Users\icosnet\Desktop\cardp\YOLOV5project\yolov5', 'custom', 'C:/Users/icosnet/Desktop/cardp/YOLOV5project/yolov5/runs/train/exp15/weights/best.pt', source='local')
the result is not really good ( bounding box and confidence ) I don't know where is the problem.

[paramkaur10]

!python3 yolov5/export.py --weights yolov5/runs/train/results_128/weights/best.pt --img 256 --include tflite

Trying to convert best.pt to tflite. But I keep getting this error.

export: data=yolov5/data/coco128.yaml, weights=['yolov5/runs/train/results_128/weights/best.pt'], imgsz=[256], batch_size=1, device=cpu, half=False, inplace=False, keras=False, optimize=False, int8=False, dynamic=False, simplify=False, opset=17, 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=['saved_model']
YOLOv5 🚀 v7.0-66-g9650f16 Python-3.9.13 torch-1.12.1 CPU

Fusing layers...
Model summary: 267 layers, 46113663 parameters, 0 gradients, 107.7 GFLOPs

PyTorch: starting from yolov5/runs/train/results_128/weights/best.pt with output shape (1, 4032, 7) (88.4 MB)
/home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages/tensorflow/python/framework/dtypes.py:585: FutureWarning: In the future np.object will be defined as the corresponding NumPy scalar. (This may have returned Python scalars in past versions.
np.object,
/home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages/tensorflow/python/framework/dtypes.py:585: FutureWarning: In the future np.object will be defined as the corresponding NumPy scalar. (This may have returned Python scalars in past versions.
np.object,
TensorFlow SavedModel: export failure ❌ 0.4s: module 'numpy' has no attribute 'object'

[loorr]

您好，您的来信我已收到祝好~   

[Alex866850]

你好，已收到你的邮件，谢谢！

[zldrobit]

@paramkaur10 I cannot reproduce the error with a colab notebook. Plz recheck your running environment (e.g. versions of YOLOv5 and other packages, and the path of the current directory).

[VirtualChrisP]

Hello, what should I do if I want to export three output torchscript models?

[Alex866850]

你好，已收到你的邮件，谢谢！

[loorr]

您好，您的来信我已收到祝好~   

[TimotheeWrightFicha]

Hello, i'm curious about --nms and thresholds options. I exported a Yolov5 model with this option but in inference it still predict bbox with conf < conf_thres.
If this isn't the goal of these options what are their purpose ?

[Alex866850]

你好，已收到你的邮件，谢谢！

[loorr]

您好，您的来信我已收到祝好~