Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

TFLite, ONNX, CoreML, TensorRT Export #251

Open
glenn-jocher opened this issue Jun 30, 2020 · 409 comments
Open

TFLite, ONNX, CoreML, TensorRT Export #251

glenn-jocher opened this issue Jun 30, 2020 · 409 comments
Assignees
Labels
documentation Improvements or additions to documentation enhancement New feature or request

Comments

@glenn-jocher
Copy link
Member

glenn-jocher commented Jun 30, 2020

📚 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. Open In Colab

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 Open In Colab. 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:

YOLOv5 OpenVINO C++ inference examples:

TensorFlow.js Web Browser Inference

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):

Status

YOLOv5 CI

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.

@glenn-jocher glenn-jocher added the documentation Improvements or additions to documentation label Jun 30, 2020
@glenn-jocher glenn-jocher pinned this issue Jun 30, 2020
@TommyZihao
Copy link
Contributor

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

@glenn-jocher glenn-jocher changed the title ONNX and Torchscript Model Export ONNX and TorchScript Model Export Jul 1, 2020
@tienhoang1094
Copy link

tienhoang1094 commented Jul 3, 2020

image
image

it only work with 5s pretrained,

@rcg12387
Copy link

rcg12387 commented Jul 3, 2020

@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
Copy link

rcg12387 commented Jul 3, 2020

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
Copy link

rcg12387 commented Jul 4, 2020

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
Copy link

Ezra-Yu commented Jul 10, 2020

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

@glenn-jocher
Copy link
Member Author

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

@ycdhqzhiai
Copy link

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

@neverrop
Copy link

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
Copy link
Contributor

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
Copy link

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
Copy link

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
Copy link
Contributor

Hi @shenglih

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

Regards

@Mayur2992
Copy link

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
Copy link
Contributor

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
Copy link

zyyang commented Jul 29, 2020

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
Copy link

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
Copy link
Contributor

@zhepherd

Please install torch=1.5.1.

@Qausain
Copy link

Qausain commented Dec 19, 2022

@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
Copy link
Contributor

@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
Copy link

fzemi commented Jan 5, 2023

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

@Qausain
Copy link

Qausain commented Jan 5, 2023

@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
Copy link

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
Copy link

!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
Copy link

loorr commented Feb 3, 2023 via email

@Alex866850
Copy link

Alex866850 commented Feb 3, 2023 via email

@zldrobit
Copy link
Contributor

zldrobit commented Feb 6, 2023

@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
Copy link

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

@Alex866850
Copy link

Alex866850 commented Feb 20, 2023 via email

@loorr
Copy link

loorr commented Feb 20, 2023 via email

@TimotheeWrightFicha
Copy link

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
Copy link

Alex866850 commented Mar 22, 2023 via email

@loorr
Copy link

loorr commented Mar 22, 2023 via email

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
documentation Improvements or additions to documentation enhancement New feature or request
Projects
None yet
Development

No branches or pull requests