141 lines
6.2 KiB
Python
141 lines
6.2 KiB
Python
# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
|
|
"""
|
|
Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit
|
|
|
|
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/
|
|
|
|
Requirements:
|
|
$ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu # CPU
|
|
$ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU
|
|
|
|
Usage:
|
|
$ python export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ...
|
|
|
|
Inference:
|
|
$ 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
|
|
|
|
TensorFlow.js:
|
|
$ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
|
|
$ npm install
|
|
$ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model
|
|
$ npm start
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import platform
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pandas as pd
|
|
import torch
|
|
from torch.utils.mobile_optimizer import optimize_for_mobile
|
|
|
|
FILE = Path(__file__).resolve()
|
|
ROOT = FILE.parents[0] # YOLOv5 root directory
|
|
if str(ROOT) not in sys.path:
|
|
sys.path.append(str(ROOT)) # add ROOT to PATH
|
|
if platform.system() != 'Windows':
|
|
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
|
|
|
|
|
from detecttracking.utils.general import (LOGGER, Profile, colorstr, file_size, get_default_args)
|
|
|
|
|
|
MACOS = platform.system() == 'Darwin' # macOS environment
|
|
|
|
|
|
class iOSModel(torch.nn.Module):
|
|
|
|
def __init__(self, model, im):
|
|
super().__init__()
|
|
b, c, h, w = im.shape # batch, channel, height, width
|
|
self.model = model
|
|
self.nc = model.nc # number of classes
|
|
if w == h:
|
|
self.normalize = 1. / w
|
|
else:
|
|
self.normalize = torch.tensor([1. / w, 1. / h, 1. / w, 1. / h]) # broadcast (slower, smaller)
|
|
# np = model(im)[0].shape[1] # number of points
|
|
# self.normalize = torch.tensor([1. / w, 1. / h, 1. / w, 1. / h]).expand(np, 4) # explicit (faster, larger)
|
|
|
|
def forward(self, x):
|
|
xywh, conf, cls = self.model(x)[0].squeeze().split((4, 1, self.nc), 1)
|
|
return cls * conf, xywh * self.normalize # confidence (3780, 80), coordinates (3780, 4)
|
|
|
|
|
|
def export_formats():
|
|
# YOLOv5 export formats
|
|
x = [
|
|
['PyTorch', '-', '.pt', True, True],
|
|
['TorchScript', 'torchscript', '.torchscript', True, True],
|
|
['ONNX', 'onnx', '.onnx', True, True],
|
|
['OpenVINO', 'openvino', '_openvino_model', True, False],
|
|
['TensorRT', 'engine', '.engine', False, True],
|
|
['CoreML', 'coreml', '.mlmodel', True, False],
|
|
['TensorFlow SavedModel', 'saved_model', '_saved_model', True, True],
|
|
['TensorFlow GraphDef', 'pb', '.pb', True, True],
|
|
['TensorFlow Lite', 'tflite', '.tflite', True, False],
|
|
['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False, False],
|
|
['TensorFlow.js', 'tfjs', '_web_model', False, False],
|
|
['PaddlePaddle', 'paddle', '_paddle_model', True, True], ]
|
|
return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'CPU', 'GPU'])
|
|
|
|
|
|
def try_export(inner_func):
|
|
# YOLOv5 export decorator, i..e @try_export
|
|
inner_args = get_default_args(inner_func)
|
|
|
|
def outer_func(*args, **kwargs):
|
|
prefix = inner_args['prefix']
|
|
try:
|
|
with Profile() as dt:
|
|
f, model = inner_func(*args, **kwargs)
|
|
LOGGER.info(f'{prefix} export success ✅ {dt.t:.1f}s, saved as {f} ({file_size(f):.1f} MB)')
|
|
return f, model
|
|
except Exception as e:
|
|
LOGGER.info(f'{prefix} export failure ❌ {dt.t:.1f}s: {e}')
|
|
return None, None
|
|
|
|
return outer_func
|
|
|
|
|
|
@try_export
|
|
def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
|
|
# YOLOv5 TorchScript model export
|
|
LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...')
|
|
f = file.with_suffix('.torchscript')
|
|
|
|
ts = torch.jit.trace(model, im, strict=False)
|
|
d = {'shape': im.shape, 'stride': int(max(model.stride)), 'names': model.names}
|
|
extra_files = {'config.txt': json.dumps(d)} # torch._C.ExtraFilesMap()
|
|
if optimize: # https://pytorch.org/tutorials/recipes/mobile_interpreter.html
|
|
optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)
|
|
else:
|
|
ts.save(str(f), _extra_files=extra_files)
|
|
return f, None
|
|
|