Yolov5: ONNX, TorchScript and CoreML Model Export

Created on 1 Jul 2020 · 124Comments · Source: ultralytics/yolov5

🚀 This guide explains how to export a trained YOLOv5 model from PyTorch to ONNX and TorchScript formats.

Before You Start

Clone this repo and install requirements.txt dependencies, including Python>=3.8 and PyTorch==1.6 and also ONNX>=1.7.

git clone https://github.com/ultralytics/yolov5  # clone repo
cd yolov5
pip install -r requirements.txt  # base requirements
pip install onnx>=1.7.0  # for ONNX export
pip install coremltools==4.0  # for CoreML export

Export a Trained YOLOv5 Model

This command exports a pretrained YOLOv5s model to ONNX, TorchScript and CoreML formats. yolov5s.pt is the lightest and fastest model available. Other options are yolov5m.pt, yolov5l.pt and yolov5x.pt, or you own checkpoint from training a custom dataset runs/exp0/weights/best.pt. For details on all available models please see our README table.

python models/export.py --weights yolov5s.pt --img 640 --batch 1  # export at 640x640 with batch size 1

Output:

Namespace(batch_size=1, img_size=[640, 640], weights='./yolov5s.pt')
Downloading https://github.com/ultralytics/yolov5/releases/download/v3.0/yolov5s.pt to ./yolov5s.pt...
100%|██████████| 14.5M/14.5M [00:02<00:00, 5.83MB/s]

Fusing layers... 
Model Summary: 140 layers, 7.45958e+06 parameters, 0 gradients, 17.5 GFLOPS

Starting TorchScript export with torch 1.6.0...
TorchScript export success, saved as ./yolov5s.torchscript.pt

Starting ONNX export with onnx 1.7.0...
ONNX export success, saved as ./yolov5s.onnx

Starting CoreML export with coremltools 4.0...
Converting graph.
...
Translating MIL ==> MLModel Ops: 100%|██████████| 1077/1077 [00:00<00:00, 1236.09 ops/s]
CoreML export success, saved as ./yolov5s.mlmodel

Export complete (16.67s). Visualize with https://github.com/lutzroeder/netron.

The 3 exported models will be saved alongside the original PyTorch model:
Screenshot 2020-10-16 at 12 51 34

Netron Viewer is highly recommended for viewing exported models:
Screenshot 2020-10-16 at 12 54 27

TensorRT Deployment

For deployment of YOLOv5 from PyTorch *.pt weights to NVIDIA TensorRT see https://github.com/wang-xinyu/tensorrtx.

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

Google Colab Notebook with free GPU:
Kaggle Notebook with free GPU: https://www.kaggle.com/ultralytics/yolov5
Google Cloud Deep Learning VM. See GCP Quickstart Guide
Docker Image https://hub.docker.com/r/ultralytics/yolov5. See Docker Quickstart Guide

documentation enhancement

Source

glenn-jocher

👍21 ❤9 🚀8 🎉6 😄3

Most helpful comment

is this possible to create with variable batch size onnx model?

                torch.onnx.export(self.model,  # model being run
                                  inputs,  # model input (or a tuple for multiple inputs)
                                  export_onnx_file,  # where to save the model (can be a file or file-like object)
                                  export_params=True,  # store the trained parameter weights inside the model file
                                  opset_version=10,  # the ONNX version to export the model to
                                  do_constant_folding=True,  # whether to execute constant folding for optimization
                                  input_names=['input'],  # the model's input names
                                  output_names=['output'],  # the model's output names
                                  dynamic_axes={'input': {0: 'batch_size'},  # variable lenght axes  # 批处理
                                                'output': {0: 'batch_size'}})

nobody-cheng on 7 Sep 2020

👍5 🎉2

All 124 comments

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

TommyZihao on 1 Jul 2020

😄1

it only work with 5s pretrained,

tienhoang1094 on 3 Jul 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 on 3 Jul 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 on 3 Jul 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.

rcg12387 on 4 Jul 2020

❤3

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

Ezra-Yu on 10 Jul 2020

👍5

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

glenn-jocher on 10 Jul 2020

😄1

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

ycdhqzhiai on 14 Jul 2020

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.

neverrop on 14 Jul 2020

👍5

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: https://github.com/ultralytics/yolov5/issues/343#issuecomment-658021043

Hope this is useful!

dlawrences on 14 Jul 2020

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｡

neverrop on 15 Jul 2020

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

shenglih on 15 Jul 2020

Hi @shenglih

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

Regards

dlawrences on 16 Jul 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?

Mayur2992 on 17 Jul 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?

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

Please see this one: https://github.com/ultralytics/yolov5/issues/315#issuecomment-656629623

dlawrences on 28 Jul 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

zyyang on 29 Jul 2020

👍2 ❤1

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

zhepherd on 29 Jul 2020

@zhepherd

Please install torch=1.5.1.

dlawrences on 29 Jul 2020

👍1

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?
Try this out:

import coremltools as ct

model = ct.converters.onnx.convert(model='my_model.onnx')

Abhimanyu8713 on 30 Jul 2020

@zhepherd

Please install torch=1.5.1.

thx it's ok

zhepherd on 31 Jul 2020

When I convert the onnx model to trt. I meet this problem:

While parsing node number 164 [Resize]:
ERROR: ModelImporter.cpp:124 In function parseGraph:
[5] Assertion failed: ctx->tensors().count(inputName)

I use tensorRT 7.0 with opset 12

VCBE123 on 3 Aug 2020

How is the output tensor meant to be read? Currently when I read the tensor it includes negative numbers and has a 5D shape. I'm also new to Yolo

BernardinD on 8 Aug 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

Yes Brother,
Thanks its working now.

Do you have any further step to deploy in ios?

Mayur2992 on 10 Aug 2020

I don't know it is okay to put my questions here. But hopefully someone could answer my questions.

I successfully converted my custom yolov5 model(train it using pre-train model yolov5x using only car, bus, truck data from 2017_train COCO datasets) and made onnx model too following instructions in this github. However, the outputs of onnx model is quite hard for me to understand and I don't know how to draw bounding boxes on original images from the outputs.

the output and my code below

"======================code======================="

layer name for onnx model

followed this, https://github.com/onnx/onnx/issues/2657

import onnx
model = onnx.load('xxx.onnx')
output =[node.name for node in model.graph.output]

input_all = [node.name for node in model.graph.input]
input_initializer = [node.name for node in model.graph.initializer]
net_feed_input = list(set(input_all) - set(input_initializer))

print('Inputs: ', net_feed_input)
print('Outputs: ', output)

intput: ['images']
output: ['output', '772', '791']

I followed this link, https://pytorch.org/docs/stable/onnx.html

import onnxruntime as ort

ort_session = ort.InferenceSession('best.onnx')

outputs = ort_session.run(None, {'images': np.random.randn(1, 3, 640, 640).astype(np.float32)})

print(outputs[0])`

"======================output========================="

Well, to put it in a nutshell, my questions below.

Could anyone tell me what it means for each dimension of the output?
Could anyone tell postpreprocessing after inferencing stages of onnx model? ( For example, https://github.com/onnx/models/blob/master/vision/object_detection_segmentation/yolov4/dependencies/inference.ipynb)

thanks

jubrowon on 12 Aug 2020

@jubrowon Follow this guy's script and the thread and you should be fine. https://github.com/ultralytics/yolov5/issues/343#issuecomment-659223637

Going through it will break down most of what you'll need in order to understand what's going on. A quick simplistic overview, by default the final postprocessing layer of the model isn't exported and that's why the ouput seems confusing

BernardinD on 12 Aug 2020

😄1

@BernardinD Thank you so much!

jubrowon on 12 Aug 2020

@glenn-jocher some notes for Windows:
it seems like setting the PYTHONPATH using set PYTHONPATH="%cd%" is not enough for torch to load the model correctly (I get an error ModuleNotFoundError: No module named 'models' from torch.load when trying to load the model). I tried a few things to make the relative import work, but couldn't find a simple solution.

What I did to make it work is to simply move export.py at the root of the project and then it exported correctly following the export command.

Ownmarc on 12 Aug 2020

@Ownmarc CI tests include export on Windows. All tests are passing. Code below, recent run here.
https://github.com/ultralytics/yolov5/blob/d2da5230533db7a2c76af1dde6d91c7e1631a1b8/.github/workflows/ci-testing.yml#L60-L75

glenn-jocher on 12 Aug 2020

@glenn-jocher ah, we have to use bash commands and not the cmd ! I tried it using bash and it worked as intended, I didn't notice I had to use bash there, I rarely use bash on Windows!

Ownmarc on 12 Aug 2020

👍1

i get an error : Can't get attribute 'Hardswish' on