|
|
|
""" |
|
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 argparse |
|
import contextlib |
|
import json |
|
import os |
|
import platform |
|
import re |
|
import subprocess |
|
import sys |
|
import time |
|
import warnings |
|
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] |
|
if str(ROOT) not in sys.path: |
|
sys.path.append(str(ROOT)) |
|
if platform.system() != "Windows": |
|
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) |
|
|
|
from models.experimental import attempt_load |
|
from models.yolo import ClassificationModel, Detect, DetectionModel, SegmentationModel |
|
from utils.dataloaders import LoadImages |
|
from utils.general import ( |
|
LOGGER, |
|
Profile, |
|
check_dataset, |
|
check_img_size, |
|
check_requirements, |
|
check_version, |
|
check_yaml, |
|
colorstr, |
|
file_size, |
|
get_default_args, |
|
print_args, |
|
url2file, |
|
yaml_save, |
|
) |
|
from utils.torch_utils import select_device, smart_inference_mode |
|
|
|
MACOS = platform.system() == "Darwin" |
|
|
|
|
|
class iOSModel(torch.nn.Module): |
|
def __init__(self, model, im): |
|
super().__init__() |
|
b, c, h, w = im.shape |
|
self.model = model |
|
self.nc = model.nc |
|
if w == h: |
|
self.normalize = 1.0 / w |
|
else: |
|
self.normalize = torch.tensor([1.0 / w, 1.0 / h, 1.0 / w, 1.0 / h]) |
|
|
|
|
|
|
|
def forward(self, x): |
|
xywh, conf, cls = self.model(x)[0].squeeze().split((4, 1, self.nc), 1) |
|
return cls * conf, xywh * self.normalize |
|
|
|
|
|
def 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): |
|
|
|
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:")): |
|
|
|
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)} |
|
if optimize: |
|
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 |
|
|
|
|
|
@try_export |
|
def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX:")): |
|
|
|
check_requirements("onnx>=1.12.0") |
|
import onnx |
|
|
|
LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__}...") |
|
f = str(file.with_suffix(".onnx")) |
|
|
|
output_names = ["output0", "output1"] if isinstance(model, SegmentationModel) else ["output0"] |
|
if dynamic: |
|
dynamic = {"images": {0: "batch", 2: "height", 3: "width"}} |
|
if isinstance(model, SegmentationModel): |
|
dynamic["output0"] = {0: "batch", 1: "anchors"} |
|
dynamic["output1"] = {0: "batch", 2: "mask_height", 3: "mask_width"} |
|
elif isinstance(model, DetectionModel): |
|
dynamic["output0"] = {0: "batch", 1: "anchors"} |
|
|
|
torch.onnx.export( |
|
model.cpu() if dynamic else model, |
|
im.cpu() if dynamic else im, |
|
f, |
|
verbose=False, |
|
opset_version=opset, |
|
do_constant_folding=True, |
|
input_names=["images"], |
|
output_names=output_names, |
|
dynamic_axes=dynamic or None, |
|
) |
|
|
|
|
|
model_onnx = onnx.load(f) |
|
onnx.checker.check_model(model_onnx) |
|
|
|
|
|
d = {"stride": int(max(model.stride)), "names": model.names} |
|
for k, v in d.items(): |
|
meta = model_onnx.metadata_props.add() |
|
meta.key, meta.value = k, str(v) |
|
onnx.save(model_onnx, f) |
|
|
|
|
|
if simplify: |
|
try: |
|
cuda = torch.cuda.is_available() |
|
check_requirements(("onnxruntime-gpu" if cuda else "onnxruntime", "onnx-simplifier>=0.4.1")) |
|
import onnxsim |
|
|
|
LOGGER.info(f"{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...") |
|
model_onnx, check = onnxsim.simplify(model_onnx) |
|
assert check, "assert check failed" |
|
onnx.save(model_onnx, f) |
|
except Exception as e: |
|
LOGGER.info(f"{prefix} simplifier failure: {e}") |
|
return f, model_onnx |
|
|
|
|
|
@try_export |
|
def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:")): |
|
|
|
check_requirements("openvino-dev>=2023.0") |
|
import openvino.runtime as ov |
|
from openvino.tools import mo |
|
|
|
LOGGER.info(f"\n{prefix} starting export with openvino {ov.__version__}...") |
|
f = str(file).replace(file.suffix, f"_{'int8_' if int8 else ''}openvino_model{os.sep}") |
|
f_onnx = file.with_suffix(".onnx") |
|
f_ov = str(Path(f) / file.with_suffix(".xml").name) |
|
|
|
ov_model = mo.convert_model(f_onnx, model_name=file.stem, framework="onnx", compress_to_fp16=half) |
|
|
|
if int8: |
|
check_requirements("nncf>=2.5.0") |
|
import nncf |
|
import numpy as np |
|
|
|
from utils.dataloaders import create_dataloader |
|
|
|
def gen_dataloader(yaml_path, task="train", imgsz=640, workers=4): |
|
data_yaml = check_yaml(yaml_path) |
|
data = check_dataset(data_yaml) |
|
dataloader = create_dataloader( |
|
data[task], imgsz=imgsz, batch_size=1, stride=32, pad=0.5, single_cls=False, rect=False, workers=workers |
|
)[0] |
|
return dataloader |
|
|
|
|
|
|
|
def transform_fn(data_item): |
|
""" |
|
Quantization transform function. |
|
|
|
Extracts and preprocess input data from dataloader item for quantization. |
|
Parameters: |
|
data_item: Tuple with data item produced by DataLoader during iteration |
|
Returns: |
|
input_tensor: Input data for quantization |
|
""" |
|
assert data_item[0].dtype == torch.uint8, "input image must be uint8 for the quantization preprocessing" |
|
|
|
img = data_item[0].numpy().astype(np.float32) |
|
img /= 255.0 |
|
return np.expand_dims(img, 0) if img.ndim == 3 else img |
|
|
|
ds = gen_dataloader(data) |
|
quantization_dataset = nncf.Dataset(ds, transform_fn) |
|
ov_model = nncf.quantize(ov_model, quantization_dataset, preset=nncf.QuantizationPreset.MIXED) |
|
|
|
ov.serialize(ov_model, f_ov) |
|
yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata) |
|
return f, None |
|
|
|
|
|
@try_export |
|
def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")): |
|
|
|
check_requirements(("paddlepaddle", "x2paddle")) |
|
import x2paddle |
|
from x2paddle.convert import pytorch2paddle |
|
|
|
LOGGER.info(f"\n{prefix} starting export with X2Paddle {x2paddle.__version__}...") |
|
f = str(file).replace(".pt", f"_paddle_model{os.sep}") |
|
|
|
pytorch2paddle(module=model, save_dir=f, jit_type="trace", input_examples=[im]) |
|
yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata) |
|
return f, None |
|
|
|
|
|
@try_export |
|
def export_coreml(model, im, file, int8, half, nms, prefix=colorstr("CoreML:")): |
|
|
|
check_requirements("coremltools") |
|
import coremltools as ct |
|
|
|
LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...") |
|
f = file.with_suffix(".mlmodel") |
|
|
|
if nms: |
|
model = iOSModel(model, im) |
|
ts = torch.jit.trace(model, im, strict=False) |
|
ct_model = ct.convert(ts, inputs=[ct.ImageType("image", shape=im.shape, scale=1 / 255, bias=[0, 0, 0])]) |
|
bits, mode = (8, "kmeans_lut") if int8 else (16, "linear") if half else (32, None) |
|
if bits < 32: |
|
if MACOS: |
|
with warnings.catch_warnings(): |
|
warnings.filterwarnings("ignore", category=DeprecationWarning) |
|
ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode) |
|
else: |
|
print(f"{prefix} quantization only supported on macOS, skipping...") |
|
ct_model.save(f) |
|
return f, ct_model |
|
|
|
|
|
@try_export |
|
def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr("TensorRT:")): |
|
|
|
assert im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. `python export.py --device 0`" |
|
try: |
|
import tensorrt as trt |
|
except Exception: |
|
if platform.system() == "Linux": |
|
check_requirements("nvidia-tensorrt", cmds="-U --index-url https://pypi.ngc.nvidia.com") |
|
import tensorrt as trt |
|
|
|
if trt.__version__[0] == "7": |
|
grid = model.model[-1].anchor_grid |
|
model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid] |
|
export_onnx(model, im, file, 12, dynamic, simplify) |
|
model.model[-1].anchor_grid = grid |
|
else: |
|
check_version(trt.__version__, "8.0.0", hard=True) |
|
export_onnx(model, im, file, 12, dynamic, simplify) |
|
onnx = file.with_suffix(".onnx") |
|
|
|
LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...") |
|
assert onnx.exists(), f"failed to export ONNX file: {onnx}" |
|
f = file.with_suffix(".engine") |
|
logger = trt.Logger(trt.Logger.INFO) |
|
if verbose: |
|
logger.min_severity = trt.Logger.Severity.VERBOSE |
|
|
|
builder = trt.Builder(logger) |
|
config = builder.create_builder_config() |
|
config.max_workspace_size = workspace * 1 << 30 |
|
|
|
|
|
flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) |
|
network = builder.create_network(flag) |
|
parser = trt.OnnxParser(network, logger) |
|
if not parser.parse_from_file(str(onnx)): |
|
raise RuntimeError(f"failed to load ONNX file: {onnx}") |
|
|
|
inputs = [network.get_input(i) for i in range(network.num_inputs)] |
|
outputs = [network.get_output(i) for i in range(network.num_outputs)] |
|
for inp in inputs: |
|
LOGGER.info(f'{prefix} input "{inp.name}" with shape{inp.shape} {inp.dtype}') |
|
for out in outputs: |
|
LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}') |
|
|
|
if dynamic: |
|
if im.shape[0] <= 1: |
|
LOGGER.warning(f"{prefix} WARNING ⚠️ --dynamic model requires maximum --batch-size argument") |
|
profile = builder.create_optimization_profile() |
|
for inp in inputs: |
|
profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape) |
|
config.add_optimization_profile(profile) |
|
|
|
LOGGER.info(f"{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}") |
|
if builder.platform_has_fast_fp16 and half: |
|
config.set_flag(trt.BuilderFlag.FP16) |
|
with builder.build_engine(network, config) as engine, open(f, "wb") as t: |
|
t.write(engine.serialize()) |
|
return f, None |
|
|
|
|
|
@try_export |
|
def export_saved_model( |
|
model, |
|
im, |
|
file, |
|
dynamic, |
|
tf_nms=False, |
|
agnostic_nms=False, |
|
topk_per_class=100, |
|
topk_all=100, |
|
iou_thres=0.45, |
|
conf_thres=0.25, |
|
keras=False, |
|
prefix=colorstr("TensorFlow SavedModel:"), |
|
): |
|
|
|
try: |
|
import tensorflow as tf |
|
except Exception: |
|
check_requirements(f"tensorflow{'' if torch.cuda.is_available() else '-macos' if MACOS else '-cpu'}") |
|
import tensorflow as tf |
|
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 |
|
|
|
from models.tf import TFModel |
|
|
|
LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") |
|
if tf.__version__ > "2.13.1": |
|
helper_url = "https://github.com/ultralytics/yolov5/issues/12489" |
|
LOGGER.info( |
|
f"WARNING ⚠️ using Tensorflow {tf.__version__} > 2.13.1 might cause issue when exporting the model to tflite {helper_url}" |
|
) |
|
f = str(file).replace(".pt", "_saved_model") |
|
batch_size, ch, *imgsz = list(im.shape) |
|
|
|
tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz) |
|
im = tf.zeros((batch_size, *imgsz, ch)) |
|
_ = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres) |
|
inputs = tf.keras.Input(shape=(*imgsz, ch), batch_size=None if dynamic else batch_size) |
|
outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres) |
|
keras_model = tf.keras.Model(inputs=inputs, outputs=outputs) |
|
keras_model.trainable = False |
|
keras_model.summary() |
|
if keras: |
|
keras_model.save(f, save_format="tf") |
|
else: |
|
spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype) |
|
m = tf.function(lambda x: keras_model(x)) |
|
m = m.get_concrete_function(spec) |
|
frozen_func = convert_variables_to_constants_v2(m) |
|
tfm = tf.Module() |
|
tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if tf_nms else frozen_func(x), [spec]) |
|
tfm.__call__(im) |
|
tf.saved_model.save( |
|
tfm, |
|
f, |
|
options=tf.saved_model.SaveOptions(experimental_custom_gradients=False) |
|
if check_version(tf.__version__, "2.6") |
|
else tf.saved_model.SaveOptions(), |
|
) |
|
return f, keras_model |
|
|
|
|
|
@try_export |
|
def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")): |
|
|
|
import tensorflow as tf |
|
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 |
|
|
|
LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") |
|
f = file.with_suffix(".pb") |
|
|
|
m = tf.function(lambda x: keras_model(x)) |
|
m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)) |
|
frozen_func = convert_variables_to_constants_v2(m) |
|
frozen_func.graph.as_graph_def() |
|
tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False) |
|
return f, None |
|
|
|
|
|
@try_export |
|
def export_tflite( |
|
keras_model, im, file, int8, per_tensor, data, nms, agnostic_nms, prefix=colorstr("TensorFlow Lite:") |
|
): |
|
|
|
import tensorflow as tf |
|
|
|
LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") |
|
batch_size, ch, *imgsz = list(im.shape) |
|
f = str(file).replace(".pt", "-fp16.tflite") |
|
|
|
converter = tf.lite.TFLiteConverter.from_keras_model(keras_model) |
|
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS] |
|
converter.target_spec.supported_types = [tf.float16] |
|
converter.optimizations = [tf.lite.Optimize.DEFAULT] |
|
if int8: |
|
from models.tf import representative_dataset_gen |
|
|
|
dataset = LoadImages(check_dataset(check_yaml(data))["train"], img_size=imgsz, auto=False) |
|
converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100) |
|
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] |
|
converter.target_spec.supported_types = [] |
|
converter.inference_input_type = tf.uint8 |
|
converter.inference_output_type = tf.uint8 |
|
converter.experimental_new_quantizer = True |
|
if per_tensor: |
|
converter._experimental_disable_per_channel = True |
|
f = str(file).replace(".pt", "-int8.tflite") |
|
if nms or agnostic_nms: |
|
converter.target_spec.supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS) |
|
|
|
tflite_model = converter.convert() |
|
open(f, "wb").write(tflite_model) |
|
return f, None |
|
|
|
|
|
@try_export |
|
def export_edgetpu(file, prefix=colorstr("Edge TPU:")): |
|
|
|
cmd = "edgetpu_compiler --version" |
|
help_url = "https://coral.ai/docs/edgetpu/compiler/" |
|
assert platform.system() == "Linux", f"export only supported on Linux. See {help_url}" |
|
if subprocess.run(f"{cmd} > /dev/null 2>&1", shell=True).returncode != 0: |
|
LOGGER.info(f"\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}") |
|
sudo = subprocess.run("sudo --version >/dev/null", shell=True).returncode == 0 |
|
for c in ( |
|
"curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -", |
|
'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list', |
|
"sudo apt-get update", |
|
"sudo apt-get install edgetpu-compiler", |
|
): |
|
subprocess.run(c if sudo else c.replace("sudo ", ""), shell=True, check=True) |
|
ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1] |
|
|
|
LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...") |
|
f = str(file).replace(".pt", "-int8_edgetpu.tflite") |
|
f_tfl = str(file).replace(".pt", "-int8.tflite") |
|
|
|
subprocess.run( |
|
[ |
|
"edgetpu_compiler", |
|
"-s", |
|
"-d", |
|
"-k", |
|
"10", |
|
"--out_dir", |
|
str(file.parent), |
|
f_tfl, |
|
], |
|
check=True, |
|
) |
|
return f, None |
|
|
|
|
|
@try_export |
|
def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")): |
|
|
|
check_requirements("tensorflowjs") |
|
import tensorflowjs as tfjs |
|
|
|
LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...") |
|
f = str(file).replace(".pt", "_web_model") |
|
f_pb = file.with_suffix(".pb") |
|
f_json = f"{f}/model.json" |
|
|
|
args = [ |
|
"tensorflowjs_converter", |
|
"--input_format=tf_frozen_model", |
|
"--quantize_uint8" if int8 else "", |
|
"--output_node_names=Identity,Identity_1,Identity_2,Identity_3", |
|
str(f_pb), |
|
f, |
|
] |
|
subprocess.run([arg for arg in args if arg], check=True) |
|
|
|
json = Path(f_json).read_text() |
|
with open(f_json, "w") as j: |
|
subst = re.sub( |
|
r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, ' |
|
r'"Identity.?.?": {"name": "Identity.?.?"}, ' |
|
r'"Identity.?.?": {"name": "Identity.?.?"}, ' |
|
r'"Identity.?.?": {"name": "Identity.?.?"}}}', |
|
r'{"outputs": {"Identity": {"name": "Identity"}, ' |
|
r'"Identity_1": {"name": "Identity_1"}, ' |
|
r'"Identity_2": {"name": "Identity_2"}, ' |
|
r'"Identity_3": {"name": "Identity_3"}}}', |
|
json, |
|
) |
|
j.write(subst) |
|
return f, None |
|
|
|
|
|
def add_tflite_metadata(file, metadata, num_outputs): |
|
|
|
with contextlib.suppress(ImportError): |
|
|
|
from tflite_support import flatbuffers |
|
from tflite_support import metadata as _metadata |
|
from tflite_support import metadata_schema_py_generated as _metadata_fb |
|
|
|
tmp_file = Path("/tmp/meta.txt") |
|
with open(tmp_file, "w") as meta_f: |
|
meta_f.write(str(metadata)) |
|
|
|
model_meta = _metadata_fb.ModelMetadataT() |
|
label_file = _metadata_fb.AssociatedFileT() |
|
label_file.name = tmp_file.name |
|
model_meta.associatedFiles = [label_file] |
|
|
|
subgraph = _metadata_fb.SubGraphMetadataT() |
|
subgraph.inputTensorMetadata = [_metadata_fb.TensorMetadataT()] |
|
subgraph.outputTensorMetadata = [_metadata_fb.TensorMetadataT()] * num_outputs |
|
model_meta.subgraphMetadata = [subgraph] |
|
|
|
b = flatbuffers.Builder(0) |
|
b.Finish(model_meta.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER) |
|
metadata_buf = b.Output() |
|
|
|
populator = _metadata.MetadataPopulator.with_model_file(file) |
|
populator.load_metadata_buffer(metadata_buf) |
|
populator.load_associated_files([str(tmp_file)]) |
|
populator.populate() |
|
tmp_file.unlink() |
|
|
|
|
|
def pipeline_coreml(model, im, file, names, y, prefix=colorstr("CoreML Pipeline:")): |
|
|
|
import coremltools as ct |
|
from PIL import Image |
|
|
|
print(f"{prefix} starting pipeline with coremltools {ct.__version__}...") |
|
batch_size, ch, h, w = list(im.shape) |
|
t = time.time() |
|
|
|
|
|
spec = model.get_spec() |
|
out0, out1 = iter(spec.description.output) |
|
if platform.system() == "Darwin": |
|
img = Image.new("RGB", (w, h)) |
|
|
|
out = model.predict({"image": img}) |
|
out0_shape, out1_shape = out[out0.name].shape, out[out1.name].shape |
|
else: |
|
s = tuple(y[0].shape) |
|
out0_shape, out1_shape = (s[1], s[2] - 5), (s[1], 4) |
|
|
|
|
|
nx, ny = spec.description.input[0].type.imageType.width, spec.description.input[0].type.imageType.height |
|
na, nc = out0_shape |
|
|
|
assert len(names) == nc, f"{len(names)} names found for nc={nc}" |
|
|
|
|
|
out0.type.multiArrayType.shape[:] = out0_shape |
|
out1.type.multiArrayType.shape[:] = out1_shape |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
print(spec.description) |
|
|
|
|
|
model = ct.models.MLModel(spec) |
|
|
|
|
|
nms_spec = ct.proto.Model_pb2.Model() |
|
nms_spec.specificationVersion = 5 |
|
for i in range(2): |
|
decoder_output = model._spec.description.output[i].SerializeToString() |
|
nms_spec.description.input.add() |
|
nms_spec.description.input[i].ParseFromString(decoder_output) |
|
nms_spec.description.output.add() |
|
nms_spec.description.output[i].ParseFromString(decoder_output) |
|
|
|
nms_spec.description.output[0].name = "confidence" |
|
nms_spec.description.output[1].name = "coordinates" |
|
|
|
output_sizes = [nc, 4] |
|
for i in range(2): |
|
ma_type = nms_spec.description.output[i].type.multiArrayType |
|
ma_type.shapeRange.sizeRanges.add() |
|
ma_type.shapeRange.sizeRanges[0].lowerBound = 0 |
|
ma_type.shapeRange.sizeRanges[0].upperBound = -1 |
|
ma_type.shapeRange.sizeRanges.add() |
|
ma_type.shapeRange.sizeRanges[1].lowerBound = output_sizes[i] |
|
ma_type.shapeRange.sizeRanges[1].upperBound = output_sizes[i] |
|
del ma_type.shape[:] |
|
|
|
nms = nms_spec.nonMaximumSuppression |
|
nms.confidenceInputFeatureName = out0.name |
|
nms.coordinatesInputFeatureName = out1.name |
|
nms.confidenceOutputFeatureName = "confidence" |
|
nms.coordinatesOutputFeatureName = "coordinates" |
|
nms.iouThresholdInputFeatureName = "iouThreshold" |
|
nms.confidenceThresholdInputFeatureName = "confidenceThreshold" |
|
nms.iouThreshold = 0.45 |
|
nms.confidenceThreshold = 0.25 |
|
nms.pickTop.perClass = True |
|
nms.stringClassLabels.vector.extend(names.values()) |
|
nms_model = ct.models.MLModel(nms_spec) |
|
|
|
|
|
pipeline = ct.models.pipeline.Pipeline( |
|
input_features=[ |
|
("image", ct.models.datatypes.Array(3, ny, nx)), |
|
("iouThreshold", ct.models.datatypes.Double()), |
|
("confidenceThreshold", ct.models.datatypes.Double()), |
|
], |
|
output_features=["confidence", "coordinates"], |
|
) |
|
pipeline.add_model(model) |
|
pipeline.add_model(nms_model) |
|
|
|
|
|
pipeline.spec.description.input[0].ParseFromString(model._spec.description.input[0].SerializeToString()) |
|
pipeline.spec.description.output[0].ParseFromString(nms_model._spec.description.output[0].SerializeToString()) |
|
pipeline.spec.description.output[1].ParseFromString(nms_model._spec.description.output[1].SerializeToString()) |
|
|
|
|
|
pipeline.spec.specificationVersion = 5 |
|
pipeline.spec.description.metadata.versionString = "https://github.com/ultralytics/yolov5" |
|
pipeline.spec.description.metadata.shortDescription = "https://github.com/ultralytics/yolov5" |
|
pipeline.spec.description.metadata.author = "glenn.jocher@ultralytics.com" |
|
pipeline.spec.description.metadata.license = "https://github.com/ultralytics/yolov5/blob/master/LICENSE" |
|
pipeline.spec.description.metadata.userDefined.update( |
|
{ |
|
"classes": ",".join(names.values()), |
|
"iou_threshold": str(nms.iouThreshold), |
|
"confidence_threshold": str(nms.confidenceThreshold), |
|
} |
|
) |
|
|
|
|
|
f = file.with_suffix(".mlmodel") |
|
model = ct.models.MLModel(pipeline.spec) |
|
model.input_description["image"] = "Input image" |
|
model.input_description["iouThreshold"] = f"(optional) IOU Threshold override (default: {nms.iouThreshold})" |
|
model.input_description[ |
|
"confidenceThreshold" |
|
] = f"(optional) Confidence Threshold override (default: {nms.confidenceThreshold})" |
|
model.output_description["confidence"] = 'Boxes × Class confidence (see user-defined metadata "classes")' |
|
model.output_description["coordinates"] = "Boxes × [x, y, width, height] (relative to image size)" |
|
model.save(f) |
|
print(f"{prefix} pipeline success ({time.time() - t:.2f}s), saved as {f} ({file_size(f):.1f} MB)") |
|
|
|
|
|
@smart_inference_mode() |
|
def run( |
|
data=ROOT / "data/coco128.yaml", |
|
weights=ROOT / "yolov5s.pt", |
|
imgsz=(640, 640), |
|
batch_size=1, |
|
device="cpu", |
|
include=("torchscript", "onnx"), |
|
half=False, |
|
inplace=False, |
|
keras=False, |
|
optimize=False, |
|
int8=False, |
|
per_tensor=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, |
|
): |
|
t = time.time() |
|
include = [x.lower() for x in include] |
|
fmts = tuple(export_formats()["Argument"][1:]) |
|
flags = [x in include for x in fmts] |
|
assert sum(flags) == len(include), f"ERROR: Invalid --include {include}, valid --include arguments are {fmts}" |
|
jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags |
|
file = Path(url2file(weights) if str(weights).startswith(("http:/", "https:/")) else weights) |
|
|
|
|
|
device = select_device(device) |
|
if half: |
|
assert device.type != "cpu" or coreml, "--half only compatible with GPU export, i.e. use --device 0" |
|
assert not dynamic, "--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both" |
|
model = attempt_load(weights, device=device, inplace=True, fuse=True) |
|
|
|
|
|
imgsz *= 2 if len(imgsz) == 1 else 1 |
|
if optimize: |
|
assert device.type == "cpu", "--optimize not compatible with cuda devices, i.e. use --device cpu" |
|
|
|
|
|
gs = int(max(model.stride)) |
|
imgsz = [check_img_size(x, gs) for x in imgsz] |
|
im = torch.zeros(batch_size, 3, *imgsz).to(device) |
|
|
|
|
|
model.eval() |
|
for k, m in model.named_modules(): |
|
if isinstance(m, Detect): |
|
m.inplace = inplace |
|
m.dynamic = dynamic |
|
m.export = True |
|
|
|
for _ in range(2): |
|
y = model(im) |
|
if half and not coreml: |
|
im, model = im.half(), model.half() |
|
shape = tuple((y[0] if isinstance(y, tuple) else y).shape) |
|
metadata = {"stride": int(max(model.stride)), "names": model.names} |
|
LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)") |
|
|
|
|
|
f = [""] * len(fmts) |
|
warnings.filterwarnings(action="ignore", category=torch.jit.TracerWarning) |
|
if jit: |
|
f[0], _ = export_torchscript(model, im, file, optimize) |
|
if engine: |
|
f[1], _ = export_engine(model, im, file, half, dynamic, simplify, workspace, verbose) |
|
if onnx or xml: |
|
f[2], _ = export_onnx(model, im, file, opset, dynamic, simplify) |
|
if xml: |
|
f[3], _ = export_openvino(file, metadata, half, int8, data) |
|
if coreml: |
|
f[4], ct_model = export_coreml(model, im, file, int8, half, nms) |
|
if nms: |
|
pipeline_coreml(ct_model, im, file, model.names, y) |
|
if any((saved_model, pb, tflite, edgetpu, tfjs)): |
|
assert not tflite or not tfjs, "TFLite and TF.js models must be exported separately, please pass only one type." |
|
assert not isinstance(model, ClassificationModel), "ClassificationModel export to TF formats not yet supported." |
|
f[5], s_model = export_saved_model( |
|
model.cpu(), |
|
im, |
|
file, |
|
dynamic, |
|
tf_nms=nms or agnostic_nms or tfjs, |
|
agnostic_nms=agnostic_nms or tfjs, |
|
topk_per_class=topk_per_class, |
|
topk_all=topk_all, |
|
iou_thres=iou_thres, |
|
conf_thres=conf_thres, |
|
keras=keras, |
|
) |
|
if pb or tfjs: |
|
f[6], _ = export_pb(s_model, file) |
|
if tflite or edgetpu: |
|
f[7], _ = export_tflite( |
|
s_model, im, file, int8 or edgetpu, per_tensor, data=data, nms=nms, agnostic_nms=agnostic_nms |
|
) |
|
if edgetpu: |
|
f[8], _ = export_edgetpu(file) |
|
add_tflite_metadata(f[8] or f[7], metadata, num_outputs=len(s_model.outputs)) |
|
if tfjs: |
|
f[9], _ = export_tfjs(file, int8) |
|
if paddle: |
|
f[10], _ = export_paddle(model, im, file, metadata) |
|
|
|
|
|
f = [str(x) for x in f if x] |
|
if any(f): |
|
cls, det, seg = (isinstance(model, x) for x in (ClassificationModel, DetectionModel, SegmentationModel)) |
|
det &= not seg |
|
dir = Path("segment" if seg else "classify" if cls else "") |
|
h = "--half" if half else "" |
|
s = ( |
|
"# WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference" |
|
if cls |
|
else "# WARNING ⚠️ SegmentationModel not yet supported for PyTorch Hub AutoShape inference" |
|
if seg |
|
else "" |
|
) |
|
LOGGER.info( |
|
f'\nExport complete ({time.time() - t:.1f}s)' |
|
f"\nResults saved to {colorstr('bold', file.parent.resolve())}" |
|
f"\nDetect: python {dir / ('detect.py' if det else 'predict.py')} --weights {f[-1]} {h}" |
|
f"\nValidate: python {dir / 'val.py'} --weights {f[-1]} {h}" |
|
f"\nPyTorch Hub: model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}') {s}" |
|
f'\nVisualize: https://netron.app' |
|
) |
|
return f |
|
|
|
|
|
def parse_opt(known=False): |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path") |
|
parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model.pt path(s)") |
|
parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640, 640], help="image (h, w)") |
|
parser.add_argument("--batch-size", type=int, default=1, help="batch size") |
|
parser.add_argument("--device", default="cpu", help="cuda device, i.e. 0 or 0,1,2,3 or cpu") |
|
parser.add_argument("--half", action="store_true", help="FP16 half-precision export") |
|
parser.add_argument("--inplace", action="store_true", help="set YOLOv5 Detect() inplace=True") |
|
parser.add_argument("--keras", action="store_true", help="TF: use Keras") |
|
parser.add_argument("--optimize", action="store_true", help="TorchScript: optimize for mobile") |
|
parser.add_argument("--int8", action="store_true", help="CoreML/TF/OpenVINO INT8 quantization") |
|
parser.add_argument("--per-tensor", action="store_true", help="TF per-tensor quantization") |
|
parser.add_argument("--dynamic", action="store_true", help="ONNX/TF/TensorRT: dynamic axes") |
|
parser.add_argument("--simplify", action="store_true", help="ONNX: simplify model") |
|
parser.add_argument("--opset", type=int, default=17, help="ONNX: opset version") |
|
parser.add_argument("--verbose", action="store_true", help="TensorRT: verbose log") |
|
parser.add_argument("--workspace", type=int, default=4, help="TensorRT: workspace size (GB)") |
|
parser.add_argument("--nms", action="store_true", help="TF: add NMS to model") |
|
parser.add_argument("--agnostic-nms", action="store_true", help="TF: add agnostic NMS to model") |
|
parser.add_argument("--topk-per-class", type=int, default=100, help="TF.js NMS: topk per class to keep") |
|
parser.add_argument("--topk-all", type=int, default=100, help="TF.js NMS: topk for all classes to keep") |
|
parser.add_argument("--iou-thres", type=float, default=0.45, help="TF.js NMS: IoU threshold") |
|
parser.add_argument("--conf-thres", type=float, default=0.25, help="TF.js NMS: confidence threshold") |
|
parser.add_argument( |
|
"--include", |
|
nargs="+", |
|
default=["torchscript"], |
|
help="torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle", |
|
) |
|
opt = parser.parse_known_args()[0] if known else parser.parse_args() |
|
print_args(vars(opt)) |
|
return opt |
|
|
|
|
|
def main(opt): |
|
for opt.weights in opt.weights if isinstance(opt.weights, list) else [opt.weights]: |
|
run(**vars(opt)) |
|
|
|
|
|
if __name__ == "__main__": |
|
opt = parse_opt() |
|
main(opt) |
|
|