|
|
|
""" |
|
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/ |
|
|
|
Usage: |
|
$ python path/to/export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ... |
|
|
|
Inference: |
|
$ python path/to/detect.py --weights yolov5s.pt # PyTorch |
|
yolov5s.torchscript # TorchScript |
|
yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn |
|
yolov5s.xml # 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 |
|
|
|
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 json |
|
import os |
|
import subprocess |
|
import sys |
|
import time |
|
from pathlib import Path |
|
|
|
import torch |
|
import torch.nn as nn |
|
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)) |
|
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) |
|
|
|
from models.common import Conv |
|
from models.experimental import attempt_load |
|
from models.yolo import Detect |
|
from utils.activations import SiLU |
|
from utils.datasets import LoadImages |
|
from utils.general import (LOGGER, check_dataset, check_img_size, check_requirements, check_version, colorstr, |
|
file_size, print_args, url2file) |
|
from utils.torch_utils import select_device |
|
|
|
|
|
def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')): |
|
|
|
try: |
|
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) |
|
|
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') |
|
except Exception as e: |
|
LOGGER.info(f'{prefix} export failure: {e}') |
|
|
|
|
|
def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')): |
|
|
|
try: |
|
check_requirements(('onnx',)) |
|
import onnx |
|
|
|
LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...') |
|
f = file.with_suffix('.onnx') |
|
|
|
torch.onnx.export(model, im, f, verbose=False, opset_version=opset, |
|
training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL, |
|
do_constant_folding=not train, |
|
input_names=['images'], |
|
output_names=['output'], |
|
dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, |
|
'output': {0: 'batch', 1: 'anchors'} |
|
} if dynamic else None) |
|
|
|
|
|
model_onnx = onnx.load(f) |
|
onnx.checker.check_model(model_onnx) |
|
|
|
|
|
|
|
if simplify: |
|
try: |
|
check_requirements(('onnx-simplifier',)) |
|
import onnxsim |
|
|
|
LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...') |
|
model_onnx, check = onnxsim.simplify( |
|
model_onnx, |
|
dynamic_input_shape=dynamic, |
|
input_shapes={'images': list(im.shape)} if dynamic else None) |
|
assert check, 'assert check failed' |
|
onnx.save(model_onnx, f) |
|
except Exception as e: |
|
LOGGER.info(f'{prefix} simplifier failure: {e}') |
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') |
|
LOGGER.info(f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'") |
|
except Exception as e: |
|
LOGGER.info(f'{prefix} export failure: {e}') |
|
|
|
|
|
def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')): |
|
|
|
try: |
|
check_requirements(('openvino-dev',)) |
|
import openvino.inference_engine as ie |
|
|
|
LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...') |
|
f = str(file).replace('.pt', '_openvino_model' + os.sep) |
|
|
|
cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f}" |
|
subprocess.check_output(cmd, shell=True) |
|
|
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') |
|
except Exception as e: |
|
LOGGER.info(f'\n{prefix} export failure: {e}') |
|
|
|
|
|
def export_coreml(model, im, file, prefix=colorstr('CoreML:')): |
|
|
|
ct_model = None |
|
try: |
|
check_requirements(('coremltools',)) |
|
import coremltools as ct |
|
|
|
LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...') |
|
f = file.with_suffix('.mlmodel') |
|
|
|
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])]) |
|
ct_model.save(f) |
|
|
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') |
|
except Exception as e: |
|
LOGGER.info(f'\n{prefix} export failure: {e}') |
|
|
|
return ct_model |
|
|
|
|
|
def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')): |
|
|
|
try: |
|
check_requirements(('tensorrt',)) |
|
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, train, False, simplify) |
|
model.model[-1].anchor_grid = grid |
|
else: |
|
check_version(trt.__version__, '8.0.0', hard=True) |
|
export_onnx(model, im, file, 13, train, False, simplify) |
|
onnx = file.with_suffix('.onnx') |
|
assert onnx.exists(), f'failed to export ONNX file: {onnx}' |
|
|
|
LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...') |
|
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)] |
|
LOGGER.info(f'{prefix} Network Description:') |
|
for inp in inputs: |
|
LOGGER.info(f'{prefix}\tinput "{inp.name}" with shape {inp.shape} and dtype {inp.dtype}') |
|
for out in outputs: |
|
LOGGER.info(f'{prefix}\toutput "{out.name}" with shape {out.shape} and dtype {out.dtype}') |
|
|
|
half &= builder.platform_has_fast_fp16 |
|
LOGGER.info(f'{prefix} building FP{16 if half else 32} engine in {f}') |
|
if half: |
|
config.set_flag(trt.BuilderFlag.FP16) |
|
with builder.build_engine(network, config) as engine, open(f, 'wb') as t: |
|
t.write(engine.serialize()) |
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') |
|
|
|
except Exception as e: |
|
LOGGER.info(f'\n{prefix} export failure: {e}') |
|
|
|
|
|
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, prefix=colorstr('TensorFlow SavedModel:')): |
|
|
|
keras_model = None |
|
try: |
|
import tensorflow as tf |
|
from tensorflow import keras |
|
|
|
from models.tf import TFDetect, TFModel |
|
|
|
LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') |
|
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, 3)) |
|
y = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres) |
|
inputs = keras.Input(shape=(*imgsz, 3), 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 = keras.Model(inputs=inputs, outputs=outputs) |
|
keras_model.trainable = False |
|
keras_model.summary() |
|
keras_model.save(f, save_format='tf') |
|
|
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') |
|
except Exception as e: |
|
LOGGER.info(f'\n{prefix} export failure: {e}') |
|
|
|
return keras_model |
|
|
|
|
|
def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')): |
|
|
|
try: |
|
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) |
|
|
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') |
|
except Exception as e: |
|
LOGGER.info(f'\n{prefix} export failure: {e}') |
|
|
|
|
|
def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('TensorFlow Lite:')): |
|
|
|
try: |
|
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(data)['train'], img_size=imgsz, auto=False) |
|
converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib) |
|
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 = False |
|
f = str(file).replace('.pt', '-int8.tflite') |
|
|
|
tflite_model = converter.convert() |
|
open(f, "wb").write(tflite_model) |
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') |
|
|
|
except Exception as e: |
|
LOGGER.info(f'\n{prefix} export failure: {e}') |
|
|
|
|
|
def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')): |
|
|
|
try: |
|
cmd = 'edgetpu_compiler --version' |
|
out = subprocess.run(cmd, shell=True, capture_output=True, check=True) |
|
ver = out.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') |
|
|
|
cmd = f"edgetpu_compiler -s {f_tfl}" |
|
subprocess.run(cmd, shell=True, check=True) |
|
|
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') |
|
except Exception as e: |
|
LOGGER.info(f'\n{prefix} export failure: {e}') |
|
|
|
|
|
def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')): |
|
|
|
try: |
|
check_requirements(('tensorflowjs',)) |
|
import re |
|
|
|
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 + '/model.json' |
|
|
|
cmd = f"tensorflowjs_converter --input_format=tf_frozen_model " \ |
|
f"--output_node_names='Identity,Identity_1,Identity_2,Identity_3' {f_pb} {f}" |
|
subprocess.run(cmd, shell=True) |
|
|
|
json = open(f_json).read() |
|
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) |
|
|
|
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') |
|
except Exception as e: |
|
LOGGER.info(f'\n{prefix} export failure: {e}') |
|
|
|
|
|
@torch.no_grad() |
|
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, |
|
train=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 |
|
): |
|
t = time.time() |
|
include = [x.lower() for x in include] |
|
tf_exports = list(x in include for x in ('saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs')) |
|
file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights) |
|
|
|
|
|
imgsz *= 2 if len(imgsz) == 1 else 1 |
|
opset = 12 if ('openvino' in include) else opset |
|
|
|
|
|
device = select_device(device) |
|
assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0' |
|
model = attempt_load(weights, map_location=device, inplace=True, fuse=True) |
|
nc, names = model.nc, model.names |
|
|
|
|
|
gs = int(max(model.stride)) |
|
imgsz = [check_img_size(x, gs) for x in imgsz] |
|
im = torch.zeros(batch_size, 3, *imgsz).to(device) |
|
|
|
|
|
if half: |
|
im, model = im.half(), model.half() |
|
model.train() if train else model.eval() |
|
for k, m in model.named_modules(): |
|
if isinstance(m, Conv): |
|
if isinstance(m.act, nn.SiLU): |
|
m.act = SiLU() |
|
elif isinstance(m, Detect): |
|
m.inplace = inplace |
|
m.onnx_dynamic = dynamic |
|
|
|
|
|
for _ in range(2): |
|
y = model(im) |
|
LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)") |
|
|
|
|
|
if 'torchscript' in include: |
|
export_torchscript(model, im, file, optimize) |
|
if 'engine' in include: |
|
export_engine(model, im, file, train, half, simplify, workspace, verbose) |
|
if ('onnx' in include) or ('openvino' in include): |
|
export_onnx(model, im, file, opset, train, dynamic, simplify) |
|
if 'openvino' in include: |
|
export_openvino(model, im, file) |
|
if 'coreml' in include: |
|
export_coreml(model, im, file) |
|
|
|
|
|
if any(tf_exports): |
|
pb, tflite, edgetpu, tfjs = tf_exports[1:] |
|
if (tflite or edgetpu) and int8: |
|
check_requirements(('flatbuffers==1.12',)) |
|
assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.' |
|
model = export_saved_model(model, 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, |
|
conf_thres=conf_thres, iou_thres=iou_thres) |
|
if pb or tfjs: |
|
export_pb(model, im, file) |
|
if tflite or edgetpu: |
|
export_tflite(model, im, file, int8=int8 or edgetpu, data=data, ncalib=100) |
|
if edgetpu: |
|
export_edgetpu(model, im, file) |
|
if tfjs: |
|
export_tfjs(model, im, file) |
|
|
|
|
|
LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)' |
|
f"\nResults saved to {colorstr('bold', file.parent.resolve())}" |
|
f'\nVisualize with https://netron.app') |
|
|
|
|
|
def parse_opt(): |
|
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('--train', action='store_true', help='model.train() mode') |
|
parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile') |
|
parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization') |
|
parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes') |
|
parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model') |
|
parser.add_argument('--opset', type=int, default=12, 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', 'onnx'], |
|
help='available formats are (torchscript, onnx, engine, coreml, saved_model, pb, tflite, tfjs)') |
|
opt = parser.parse_args() |
|
print_args(FILE.stem, 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) |
|
|