k20hcmus/FishEye8K
3
1import argparse2import contextlib3import json4import os5import platform6import re7import subprocess8import sys9import time10import warnings11from pathlib import Path12 13import pandas as pd14import torch15from torch.utils.mobile_optimizer import optimize_for_mobile16 17FILE = Path(__file__).resolve()18ROOT = FILE.parents[0] # YOLO root directory19if str(ROOT) not in sys.path:20 sys.path.append(str(ROOT)) # add ROOT to PATH21if platform.system() != 'Windows':22 ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative23 24from models.experimental import attempt_load, End2End25from models.yolo import ClassificationModel, Detect, DDetect, DualDetect, DualDDetect, IDualDDetect, DetectionModel, SegmentationModel26from utils.dataloaders import LoadImages27from utils.general import (LOGGER, Profile, check_dataset, check_img_size, check_requirements, check_version,28 check_yaml, colorstr, file_size, get_default_args, print_args, url2file, yaml_save)29from utils.torch_utils import select_device, smart_inference_mode30 31MACOS = platform.system() == 'Darwin' # macOS environment32 33 34def export_formats():35 # YOLO export formats36 x = [37 ['PyTorch', '-', '.pt', True, True],38 ['TorchScript', 'torchscript', '.torchscript', True, True],39 ['ONNX', 'onnx', '.onnx', True, True],40 ['ONNX END2END', 'onnx_end2end', '_end2end.onnx', True, True],41 ['OpenVINO', 'openvino', '_openvino_model', True, False],42 ['TensorRT', 'engine', '.engine', False, True],43 ['CoreML', 'coreml', '.mlmodel', True, False],44 ['TensorFlow SavedModel', 'saved_model', '_saved_model', True, True],45 ['TensorFlow GraphDef', 'pb', '.pb', True, True],46 ['TensorFlow Lite', 'tflite', '.tflite', True, False],47 ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False, False],48 ['TensorFlow.js', 'tfjs', '_web_model', False, False],49 ['PaddlePaddle', 'paddle', '_paddle_model', True, True],]50 return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'CPU', 'GPU'])51 52 53def try_export(inner_func):54 # YOLO export decorator, i..e @try_export55 inner_args = get_default_args(inner_func)56 57 def outer_func(*args, **kwargs):58 prefix = inner_args['prefix']59 try:60 with Profile() as dt:61 f, model = inner_func(*args, **kwargs)62 LOGGER.info(f'{prefix} export success ✅ {dt.t:.1f}s, saved as {f} ({file_size(f):.1f} MB)')63 return f, model64 except Exception as e:65 LOGGER.info(f'{prefix} export failure ❌ {dt.t:.1f}s: {e}')66 return None, None67 68 return outer_func69 70 71@try_export72def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):73 # YOLO TorchScript model export74 LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...')75 f = file.with_suffix('.torchscript')76 77 ts = torch.jit.trace(model, im, strict=False)78 d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}79 extra_files = {'config.txt': json.dumps(d)} # torch._C.ExtraFilesMap()80 if optimize: # https://pytorch.org/tutorials/recipes/mobile_interpreter.html81 optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)82 else:83 ts.save(str(f), _extra_files=extra_files)84 return f, None85 86 87@try_export88def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX:')):89 # YOLO ONNX export90 check_requirements('onnx')91 import onnx92 93 LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')94 f = file.with_suffix('.onnx')95 96 output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output0']97 if dynamic:98 dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640)99 if isinstance(model, SegmentationModel):100 dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85)101 dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)102 elif isinstance(model, DetectionModel):103 dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85)104 105 torch.onnx.export(106 model.cpu() if dynamic else model, # --dynamic only compatible with cpu107 im.cpu() if dynamic else im,108 f,109 verbose=False,110 opset_version=opset,111 do_constant_folding=True,112 input_names=['images'],113 output_names=output_names,114 dynamic_axes=dynamic or None)115 116 # Checks117 model_onnx = onnx.load(f) # load onnx model118 onnx.checker.check_model(model_onnx) # check onnx model119 120 # Metadata121 d = {'stride': int(max(model.stride)), 'names': model.names}122 for k, v in d.items():123 meta = model_onnx.metadata_props.add()124 meta.key, meta.value = k, str(v)125 onnx.save(model_onnx, f)126 127 # Simplify128 if simplify:129 try:130 cuda = torch.cuda.is_available()131 check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1'))132 import onnxsim133 134 LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')135 model_onnx, check = onnxsim.simplify(model_onnx)136 assert check, 'assert check failed'137 onnx.save(model_onnx, f)138 except Exception as e:139 LOGGER.info(f'{prefix} simplifier failure: {e}')140 return f, model_onnx141 142 143@try_export144def export_onnx_end2end(model, im, file, simplify, topk_all, iou_thres, conf_thres, device, labels, prefix=colorstr('ONNX END2END:')):145 # YOLO ONNX export146 check_requirements('onnx')147 import onnx148 LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')149 f = os.path.splitext(file)[0] + "-end2end.onnx"150 batch_size = 'batch'151 152 dynamic_axes = {'images': {0 : 'batch', 2: 'height', 3:'width'}, } # variable length axes153 154 output_axes = {155 'num_dets': {0: 'batch'},156 'det_boxes': {0: 'batch'},157 'det_scores': {0: 'batch'},158 'det_classes': {0: 'batch'},159 }160 dynamic_axes.update(output_axes)161 model = End2End(model, topk_all, iou_thres, conf_thres, None ,device, labels)162 163 output_names = ['num_dets', 'det_boxes', 'det_scores', 'det_classes']164 shapes = [ batch_size, 1, batch_size, topk_all, 4,165 batch_size, topk_all, batch_size, topk_all]166 167 torch.onnx.export(model, 168 im, 169 f, 170 verbose=False, 171 export_params=True, # store the trained parameter weights inside the model file172 opset_version=12, 173 do_constant_folding=True, # whether to execute constant folding for optimization174 input_names=['images'],175 output_names=output_names,176 dynamic_axes=dynamic_axes)177 178 # Checks179 model_onnx = onnx.load(f) # load onnx model180 onnx.checker.check_model(model_onnx) # check onnx model181 for i in model_onnx.graph.output:182 for j in i.type.tensor_type.shape.dim:183 j.dim_param = str(shapes.pop(0))184 185 if simplify:186 try:187 import onnxsim188 189 print('\nStarting to simplify ONNX...')190 model_onnx, check = onnxsim.simplify(model_onnx)191 assert check, 'assert check failed'192 except Exception as e:193 print(f'Simplifier failure: {e}')194 195 # print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model196 onnx.save(model_onnx,f)197 print('ONNX export success, saved as %s' % f)198 return f, model_onnx199 200 201@try_export202def export_openvino(file, metadata, half, prefix=colorstr('OpenVINO:')):203 # YOLO OpenVINO export204 check_requirements('openvino-dev') # requires openvino-dev: https://pypi.org/project/openvino-dev/205 import openvino.inference_engine as ie206 207 LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...')208 f = str(file).replace('.pt', f'_openvino_model{os.sep}')209 210 #cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} --data_type {'FP16' if half else 'FP32'}"211 #cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} {"--compress_to_fp16" if half else ""}"212 half_arg = "--compress_to_fp16" if half else ""213 cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} {half_arg}"214 subprocess.run(cmd.split(), check=True, env=os.environ) # export215 yaml_save(Path(f) / file.with_suffix('.yaml').name, metadata) # add metadata.yaml216 return f, None217 218 219@try_export220def export_paddle(model, im, file, metadata, prefix=colorstr('PaddlePaddle:')):221 # YOLO Paddle export222 check_requirements(('paddlepaddle', 'x2paddle'))223 import x2paddle224 from x2paddle.convert import pytorch2paddle225 226 LOGGER.info(f'\n{prefix} starting export with X2Paddle {x2paddle.__version__}...')227 f = str(file).replace('.pt', f'_paddle_model{os.sep}')228 229 pytorch2paddle(module=model, save_dir=f, jit_type='trace', input_examples=[im]) # export230 yaml_save(Path(f) / file.with_suffix('.yaml').name, metadata) # add metadata.yaml231 return f, None232 233 234@try_export235def export_coreml(model, im, file, int8, half, prefix=colorstr('CoreML:')):236 # YOLO CoreML export237 check_requirements('coremltools')238 import coremltools as ct239 240 LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...')241 f = file.with_suffix('.mlmodel')242 243 ts = torch.jit.trace(model, im, strict=False) # TorchScript model244 ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])245 bits, mode = (8, 'kmeans_lut') if int8 else (16, 'linear') if half else (32, None)246 if bits < 32:247 if MACOS: # quantization only supported on macOS248 with warnings.catch_warnings():249 warnings.filterwarnings("ignore", category=DeprecationWarning) # suppress numpy==1.20 float warning250 ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)251 else:252 print(f'{prefix} quantization only supported on macOS, skipping...')253 ct_model.save(f)254 return f, ct_model255 256 257@try_export258def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):259 # YOLO TensorRT export https://developer.nvidia.com/tensorrt260 assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`'261 try:262 import tensorrt as trt263 except Exception:264 if platform.system() == 'Linux':265 check_requirements('nvidia-tensorrt', cmds='-U --index-url https://pypi.ngc.nvidia.com')266 import tensorrt as trt267 268 if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012269 grid = model.model[-1].anchor_grid270 model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]271 export_onnx(model, im, file, 12, dynamic, simplify) # opset 12272 model.model[-1].anchor_grid = grid273 else: # TensorRT >= 8274 check_version(trt.__version__, '8.0.0', hard=True) # require tensorrt>=8.0.0275 export_onnx(model, im, file, 12, dynamic, simplify) # opset 12276 onnx = file.with_suffix('.onnx')277 278 LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')279 assert onnx.exists(), f'failed to export ONNX file: {onnx}'280 f = file.with_suffix('.engine') # TensorRT engine file281 logger = trt.Logger(trt.Logger.INFO)282 if verbose:283 logger.min_severity = trt.Logger.Severity.VERBOSE284 285 builder = trt.Builder(logger)286 config = builder.create_builder_config()287 config.max_workspace_size = workspace * 1 << 30288 # config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30) # fix TRT 8.4 deprecation notice289 290 flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))291 network = builder.create_network(flag)292 parser = trt.OnnxParser(network, logger)293 if not parser.parse_from_file(str(onnx)):294 raise RuntimeError(f'failed to load ONNX file: {onnx}')295 296 inputs = [network.get_input(i) for i in range(network.num_inputs)]297 outputs = [network.get_output(i) for i in range(network.num_outputs)]298 for inp in inputs:299 LOGGER.info(f'{prefix} input "{inp.name}" with shape{inp.shape} {inp.dtype}')300 for out in outputs:301 LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}')302 303 if dynamic:304 if im.shape[0] <= 1:305 LOGGER.warning(f"{prefix} WARNING ⚠️ --dynamic model requires maximum --batch-size argument")306 profile = builder.create_optimization_profile()307 for inp in inputs:308 profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape)309 config.add_optimization_profile(profile)310 311 LOGGER.info(f'{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}')312 if builder.platform_has_fast_fp16 and half:313 config.set_flag(trt.BuilderFlag.FP16)314 with builder.build_engine(network, config) as engine, open(f, 'wb') as t:315 t.write(engine.serialize())316 return f, None317 318 319@try_export320def export_saved_model(model,321 im,322 file,323 dynamic,324 tf_nms=False,325 agnostic_nms=False,326 topk_per_class=100,327 topk_all=100,328 iou_thres=0.45,329 conf_thres=0.25,330 keras=False,331 prefix=colorstr('TensorFlow SavedModel:')):332 # YOLO TensorFlow SavedModel export333 try:334 import tensorflow as tf335 except Exception:336 check_requirements(f"tensorflow{'' if torch.cuda.is_available() else '-macos' if MACOS else '-cpu'}")337 import tensorflow as tf338 from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2339 340 from models.tf import TFModel341 342 LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')343 f = str(file).replace('.pt', '_saved_model')344 batch_size, ch, *imgsz = list(im.shape) # BCHW345 346 tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)347 im = tf.zeros((batch_size, *imgsz, ch)) # BHWC order for TensorFlow348 _ = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)349 inputs = tf.keras.Input(shape=(*imgsz, ch), batch_size=None if dynamic else batch_size)350 outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)351 keras_model = tf.keras.Model(inputs=inputs, outputs=outputs)352 keras_model.trainable = False353 keras_model.summary()354 if keras:355 keras_model.save(f, save_format='tf')356 else:357 spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)358 m = tf.function(lambda x: keras_model(x)) # full model359 m = m.get_concrete_function(spec)360 frozen_func = convert_variables_to_constants_v2(m)361 tfm = tf.Module()362 tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if tf_nms else frozen_func(x), [spec])363 tfm.__call__(im)364 tf.saved_model.save(tfm,365 f,366 options=tf.saved_model.SaveOptions(experimental_custom_gradients=False) if check_version(367 tf.__version__, '2.6') else tf.saved_model.SaveOptions())368 return f, keras_model369 370 371@try_export372def export_pb(keras_model, file, prefix=colorstr('TensorFlow GraphDef:')):373 # YOLO TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow374 import tensorflow as tf375 from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2376 377 LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')378 f = file.with_suffix('.pb')379 380 m = tf.function(lambda x: keras_model(x)) # full model381 m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))382 frozen_func = convert_variables_to_constants_v2(m)383 frozen_func.graph.as_graph_def()384 tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)385 return f, None386 387 388@try_export389def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=colorstr('TensorFlow Lite:')):390 # YOLOv5 TensorFlow Lite export391 import tensorflow as tf392 393 LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')394 batch_size, ch, *imgsz = list(im.shape) # BCHW395 f = str(file).replace('.pt', '-fp16.tflite')396 397 converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)398 converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]399 converter.target_spec.supported_types = [tf.float16]400 converter.optimizations = [tf.lite.Optimize.DEFAULT]401 if int8:402 from models.tf import representative_dataset_gen403 dataset = LoadImages(check_dataset(check_yaml(data))['train'], img_size=imgsz, auto=False)404 converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100)405 converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]406 converter.target_spec.supported_types = []407 converter.inference_input_type = tf.uint8 # or tf.int8408 converter.inference_output_type = tf.uint8 # or tf.int8409 converter.experimental_new_quantizer = True410 f = str(file).replace('.pt', '-int8.tflite')411 if nms or agnostic_nms:412 converter.target_spec.supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS)413 414 tflite_model = converter.convert()415 open(f, "wb").write(tflite_model)416 return f, None417 418 419@try_export420def export_edgetpu(file, prefix=colorstr('Edge TPU:')):421 # YOLO Edge TPU export https://coral.ai/docs/edgetpu/models-intro/422 cmd = 'edgetpu_compiler --version'423 help_url = 'https://coral.ai/docs/edgetpu/compiler/'424 assert platform.system() == 'Linux', f'export only supported on Linux. See {help_url}'425 if subprocess.run(f'{cmd} >/dev/null', shell=True).returncode != 0:426 LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}')427 sudo = subprocess.run('sudo --version >/dev/null', shell=True).returncode == 0 # sudo installed on system428 for c in (429 'curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -',430 'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list',431 'sudo apt-get update', 'sudo apt-get install edgetpu-compiler'):432 subprocess.run(c if sudo else c.replace('sudo ', ''), shell=True, check=True)433 ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]434 435 LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...')436 f = str(file).replace('.pt', '-int8_edgetpu.tflite') # Edge TPU model437 f_tfl = str(file).replace('.pt', '-int8.tflite') # TFLite model438 439 cmd = f"edgetpu_compiler -s -d -k 10 --out_dir {file.parent} {f_tfl}"440 subprocess.run(cmd.split(), check=True)441 return f, None442 443 444@try_export445def export_tfjs(file, prefix=colorstr('TensorFlow.js:')):446 # YOLO TensorFlow.js export447 check_requirements('tensorflowjs')448 import tensorflowjs as tfjs449 450 LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...')451 f = str(file).replace('.pt', '_web_model') # js dir452 f_pb = file.with_suffix('.pb') # *.pb path453 f_json = f'{f}/model.json' # *.json path454 455 cmd = f'tensorflowjs_converter --input_format=tf_frozen_model ' \456 f'--output_node_names=Identity,Identity_1,Identity_2,Identity_3 {f_pb} {f}'457 subprocess.run(cmd.split())458 459 json = Path(f_json).read_text()460 with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order461 subst = re.sub(462 r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '463 r'"Identity.?.?": {"name": "Identity.?.?"}, '464 r'"Identity.?.?": {"name": "Identity.?.?"}, '465 r'"Identity.?.?": {"name": "Identity.?.?"}}}', r'{"outputs": {"Identity": {"name": "Identity"}, '466 r'"Identity_1": {"name": "Identity_1"}, '467 r'"Identity_2": {"name": "Identity_2"}, '468 r'"Identity_3": {"name": "Identity_3"}}}', json)469 j.write(subst)470 return f, None471 472 473def add_tflite_metadata(file, metadata, num_outputs):474 # Add metadata to *.tflite models per https://www.tensorflow.org/lite/models/convert/metadata475 with contextlib.suppress(ImportError):476 # check_requirements('tflite_support')477 from tflite_support import flatbuffers478 from tflite_support import metadata as _metadata479 from tflite_support import metadata_schema_py_generated as _metadata_fb480 481 tmp_file = Path('/tmp/meta.txt')482 with open(tmp_file, 'w') as meta_f:483 meta_f.write(str(metadata))484 485 model_meta = _metadata_fb.ModelMetadataT()486 label_file = _metadata_fb.AssociatedFileT()487 label_file.name = tmp_file.name488 model_meta.associatedFiles = [label_file]489 490 subgraph = _metadata_fb.SubGraphMetadataT()491 subgraph.inputTensorMetadata = [_metadata_fb.TensorMetadataT()]492 subgraph.outputTensorMetadata = [_metadata_fb.TensorMetadataT()] * num_outputs493 model_meta.subgraphMetadata = [subgraph]494 495 b = flatbuffers.Builder(0)496 b.Finish(model_meta.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)497 metadata_buf = b.Output()498 499 populator = _metadata.MetadataPopulator.with_model_file(file)500 populator.load_metadata_buffer(metadata_buf)501 populator.load_associated_files([str(tmp_file)])502 populator.populate()503 tmp_file.unlink()504 505 506@smart_inference_mode()507def run(508 data=ROOT / 'data/coco.yaml', # 'dataset.yaml path'509 weights=ROOT / 'yolo.pt', # weights path510 imgsz=(640, 640), # image (height, width)511 batch_size=1, # batch size512 device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu513 include=('torchscript', 'onnx'), # include formats514 half=False, # FP16 half-precision export515 inplace=False, # set YOLO Detect() inplace=True516 keras=False, # use Keras517 optimize=False, # TorchScript: optimize for mobile518 int8=False, # CoreML/TF INT8 quantization519 dynamic=False, # ONNX/TF/TensorRT: dynamic axes520 simplify=False, # ONNX: simplify model521 opset=12, # ONNX: opset version522 verbose=False, # TensorRT: verbose log523 workspace=4, # TensorRT: workspace size (GB)524 nms=False, # TF: add NMS to model525 agnostic_nms=False, # TF: add agnostic NMS to model526 topk_per_class=100, # TF.js NMS: topk per class to keep527 topk_all=100, # TF.js NMS: topk for all classes to keep528 iou_thres=0.45, # TF.js NMS: IoU threshold529 conf_thres=0.25, # TF.js NMS: confidence threshold530):531 t = time.time()532 include = [x.lower() for x in include] # to lowercase533 fmts = tuple(export_formats()['Argument'][1:]) # --include arguments534 flags = [x in include for x in fmts]535 assert sum(flags) == len(include), f'ERROR: Invalid --include {include}, valid --include arguments are {fmts}'536 jit, onnx, onnx_end2end, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags # export booleans537 file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights) # PyTorch weights538 539 # Load PyTorch model540 device = select_device(device)541 if half:542 assert device.type != 'cpu' or coreml, '--half only compatible with GPU export, i.e. use --device 0'543 assert not dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both'544 model = attempt_load(weights, device=device, inplace=True, fuse=True) # load FP32 model545 546 # Checks547 imgsz *= 2 if len(imgsz) == 1 else 1 # expand548 if optimize:549 assert device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu'550 551 # Input552 gs = int(max(model.stride)) # grid size (max stride)553 imgsz = [check_img_size(x, gs) for x in imgsz] # verify img_size are gs-multiples554 im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection555 556 # Update model557 model.eval()558 for k, m in model.named_modules():559 if isinstance(m, (Detect, DDetect, DualDetect, DualDDetect, IDualDDetect)):560 m.inplace = inplace561 m.dynamic = dynamic562 m.export = True563 564 for _ in range(2):565 y = model(im) # dry runs566 if half and not coreml:567 im, model = im.half(), model.half() # to FP16568 shape = tuple((y[0] if isinstance(y, (tuple, list)) else y).shape) # model output shape569 metadata = {'stride': int(max(model.stride)), 'names': model.names} # model metadata570 LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)")571 572 # Exports573 f = [''] * len(fmts) # exported filenames574 warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) # suppress TracerWarning575 if jit: # TorchScript576 f[0], _ = export_torchscript(model, im, file, optimize)577 if engine: # TensorRT required before ONNX578 f[1], _ = export_engine(model, im, file, half, dynamic, simplify, workspace, verbose)579 if onnx or xml: # OpenVINO requires ONNX580 f[2], _ = export_onnx(model, im, file, opset, dynamic, simplify)581 if onnx_end2end:582 if isinstance(model, DetectionModel):583 labels = model.names584 f[2], _ = export_onnx_end2end(model, im, file, simplify, topk_all, iou_thres, conf_thres, device, len(labels))585 else:586 raise RuntimeError("The model is not a DetectionModel.")587 if xml: # OpenVINO588 f[3], _ = export_openvino(file, metadata, half)589 if coreml: # CoreML590 f[4], _ = export_coreml(model, im, file, int8, half)591 if any((saved_model, pb, tflite, edgetpu, tfjs)): # TensorFlow formats592 assert not tflite or not tfjs, 'TFLite and TF.js models must be exported separately, please pass only one type.'593 assert not isinstance(model, ClassificationModel), 'ClassificationModel export to TF formats not yet supported.'594 f[5], s_model = export_saved_model(model.cpu(),595 im,596 file,597 dynamic,598 tf_nms=nms or agnostic_nms or tfjs,599 agnostic_nms=agnostic_nms or tfjs,600 topk_per_class=topk_per_class,601 topk_all=topk_all,602 iou_thres=iou_thres,603 conf_thres=conf_thres,604 keras=keras)605 if pb or tfjs: # pb prerequisite to tfjs606 f[6], _ = export_pb(s_model, file)607 if tflite or edgetpu:608 f[7], _ = export_tflite(s_model, im, file, int8 or edgetpu, data=data, nms=nms, agnostic_nms=agnostic_nms)609 if edgetpu:610 f[8], _ = export_edgetpu(file)611 add_tflite_metadata(f[8] or f[7], metadata, num_outputs=len(s_model.outputs))612 if tfjs:613 f[9], _ = export_tfjs(file)614 if paddle: # PaddlePaddle615 f[10], _ = export_paddle(model, im, file, metadata)616 617 # Finish618 f = [str(x) for x in f if x] # filter out '' and None619 if any(f):620 cls, det, seg = (isinstance(model, x) for x in (ClassificationModel, DetectionModel, SegmentationModel)) # type621 dir = Path('segment' if seg else 'classify' if cls else '')622 h = '--half' if half else '' # --half FP16 inference arg623 s = "# WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference" if cls else \624 "# WARNING ⚠️ SegmentationModel not yet supported for PyTorch Hub AutoShape inference" if seg else ''625 if onnx_end2end:626 LOGGER.info(f'\nExport complete ({time.time() - t:.1f}s)'627 f"\nResults saved to {colorstr('bold', file.parent.resolve())}"628 f"\nVisualize: https://netron.app")629 else:630 LOGGER.info(f'\nExport complete ({time.time() - t:.1f}s)'631 f"\nResults saved to {colorstr('bold', file.parent.resolve())}"632 f"\nDetect: python {dir / ('detect.py' if det else 'predict.py')} --weights {f[-1]} {h}"633 f"\nValidate: python {dir / 'val.py'} --weights {f[-1]} {h}"634 f"\nPyTorch Hub: model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}') {s}"635 f"\nVisualize: https://netron.app")636 return f # return list of exported files/dirs637 638 639def parse_opt():640 parser = argparse.ArgumentParser()641 parser.add_argument('--data', type=str, default=ROOT / 'data/coco.yaml', help='dataset.yaml path')642 parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolo.pt', help='model.pt path(s)')643 parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')644 parser.add_argument('--batch-size', type=int, default=1, help='batch size')645 parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')646 parser.add_argument('--half', action='store_true', help='FP16 half-precision export')647 parser.add_argument('--inplace', action='store_true', help='set YOLO Detect() inplace=True')648 parser.add_argument('--keras', action='store_true', help='TF: use Keras')649 parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')650 parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')651 parser.add_argument('--dynamic', action='store_true', help='ONNX/TF/TensorRT: dynamic axes')652 parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')653 parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version')654 parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')655 parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)')656 parser.add_argument('--nms', action='store_true', help='TF: add NMS to model')657 parser.add_argument('--agnostic-nms', action='store_true', help='TF: add agnostic NMS to model')658 parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep')659 parser.add_argument('--topk-all', type=int, default=100, help='ONNX END2END/TF.js NMS: topk for all classes to keep')660 parser.add_argument('--iou-thres', type=float, default=0.45, help='ONNX END2END/TF.js NMS: IoU threshold')661 parser.add_argument('--conf-thres', type=float, default=0.25, help='ONNX END2END/TF.js NMS: confidence threshold')662 parser.add_argument(663 '--include',664 nargs='+',665 default=['torchscript'],666 help='torchscript, onnx, onnx_end2end, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle')667 opt = parser.parse_args()668 669 if 'onnx_end2end' in opt.include: 670 opt.simplify = True671 opt.dynamic = True672 opt.inplace = True673 opt.half = False674 675 print_args(vars(opt))676 return opt677 678 679def main(opt):680 for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]):681 run(**vars(opt))682 683 684if __name__ == "__main__":685 opt = parse_opt()686 main(opt)687