From 076d73cfaab6ae536368f74f85b2f99be8be6249 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 30 Dec 2022 01:28:41 +0100 Subject: [PATCH] Create Exporter() Class (#117) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- README.md | 7 +- ultralytics/nn/tasks.py | 3 +- ultralytics/yolo/cli.py | 4 +- ultralytics/yolo/configs/default.yaml | 3 +- ultralytics/yolo/engine/exporter.py | 1045 ++++++++++--------- ultralytics/yolo/engine/model.py | 37 +- ultralytics/yolo/engine/predictor.py | 8 +- ultralytics/yolo/utils/__init__.py | 20 +- ultralytics/yolo/utils/callbacks/clearml.py | 4 +- ultralytics/yolo/utils/callbacks/wb.py | 4 +- 10 files changed, 563 insertions(+), 572 deletions(-) diff --git a/README.md b/README.md index fea7133..17148b2 100644 --- a/README.md +++ b/README.md @@ -16,9 +16,10 @@ pip install -e . ### 1. CLI To simply use the latest Ultralytics YOLO models ```bash -yolo task=detect mode=train model=yolov8n.yaml ... - classify predict yolov8n-cls.yaml - segment val yolov8n-seg.yaml +yolo task=detect mode=train model=yolov8n.yaml args=... + classify predict yolov8n-cls.yaml args=... + segment val yolov8n-seg.yaml args=... + export yolov8n.pt format=onnx ``` ### 2. Python SDK To use pythonic interface of Ultralytics YOLO model diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py index 62aef70..d7162a9 100644 --- a/ultralytics/nn/tasks.py +++ b/ultralytics/nn/tasks.py @@ -11,6 +11,7 @@ from ultralytics.nn.modules import (C1, C2, C3, C3TR, SPP, SPPF, Bottleneck, Bot Concat, Conv, ConvTranspose, Detect, DWConv, DWConvTranspose2d, Ensemble, Focus, GhostBottleneck, GhostConv, Segment) from ultralytics.yolo.utils import LOGGER, colorstr +from ultralytics.yolo.utils.checks import check_yaml from ultralytics.yolo.utils.files import yaml_load from ultralytics.yolo.utils.torch_utils import (fuse_conv_and_bn, initialize_weights, intersect_state_dicts, make_divisible, model_info, scale_img, time_sync) @@ -80,7 +81,7 @@ class DetectionModel(BaseModel): # YOLOv5 detection model def __init__(self, cfg='yolov8n.yaml', ch=3, nc=None, verbose=True): # model, input channels, number of classes super().__init__() - self.yaml = cfg if isinstance(cfg, dict) else yaml_load(cfg) # cfg dict + self.yaml = cfg if isinstance(cfg, dict) else yaml_load(check_yaml(cfg)) # cfg dict # Define model ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels diff --git a/ultralytics/yolo/cli.py b/ultralytics/yolo/cli.py index 3ca907b..4f883a5 100644 --- a/ultralytics/yolo/cli.py +++ b/ultralytics/yolo/cli.py @@ -31,7 +31,7 @@ def cli(cfg): elif task == "classify": module = yolo.v8.classify elif task == "export": - func = yolo.trainer.exporter.export_model + func = yolo.engine.exporter.export else: raise SyntaxError("task not recognized. Choices are `'detect', 'segment', 'classify'`") @@ -42,7 +42,7 @@ def cli(cfg): elif mode == "predict": func = module.predict elif mode == "export": - func = yolo.trainer.exporter.export_model + func = yolo.engine.exporter.export else: raise SyntaxError("mode not recognized. Choices are `'train', 'val', 'predict', 'export'`") func(cfg) diff --git a/ultralytics/yolo/configs/default.yaml b/ultralytics/yolo/configs/default.yaml index 9d24a09..82a7af3 100644 --- a/ultralytics/yolo/configs/default.yaml +++ b/ultralytics/yolo/configs/default.yaml @@ -29,12 +29,12 @@ image_weights: False # use weighted image selection for training rect: False # support rectangular training cos_lr: False # use cosine LR scheduler close_mosaic: 10 # disable mosaic for final 10 epochs +resume: False # Segmentation overlap_mask: True # masks overlap mask_ratio: 4 # mask downsample ratio # Classification dropout: False # use dropout -resume: False # Val/Test settings ---------------------------------------------------------------------------------------------------- @@ -65,6 +65,7 @@ agnostic_nms: False # class-agnostic NMS retina_masks: False # Export settings ------------------------------------------------------------------------------------------------------ +format: torchscript keras: False # use Keras optimize: False # TorchScript: optimize for mobile int8: False # CoreML/TF INT8 quantization diff --git a/ultralytics/yolo/engine/exporter.py b/ultralytics/yolo/engine/exporter.py index d5dd82a..7585fcd 100644 --- a/ultralytics/yolo/engine/exporter.py +++ b/ultralytics/yolo/engine/exporter.py @@ -2,51 +2,51 @@ """ Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit -Format | `export.py --include` | Model +Format | `format=argument` | Model --- | --- | --- PyTorch | - | yolov8n.pt TorchScript | `torchscript` | yolov8n.torchscript ONNX | `onnx` | yolov8n.onnx -OpenVINO | `openvino` | yolov5s_openvino_model/ +OpenVINO | `openvino` | yolov8n_openvino_model/ TensorRT | `engine` | yolov8n.engine CoreML | `coreml` | yolov8n.mlmodel -TensorFlow SavedModel | `saved_model` | yolov5s_saved_model/ +TensorFlow SavedModel | `saved_model` | yolov8n_saved_model/ TensorFlow GraphDef | `pb` | yolov8n.pb TensorFlow Lite | `tflite` | yolov8n.tflite -TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite -TensorFlow.js | `tfjs` | yolov5s_web_model/ -PaddlePaddle | `paddle` | yolov5s_paddle_model/ +TensorFlow Edge TPU | `edgetpu` | yolov8n_edgetpu.tflite +TensorFlow.js | `tfjs` | yolov8n_web_model/ +PaddlePaddle | `paddle` | yolov8n_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 yolov8n.pt --include torchscript onnx openvino engine coreml tflite ... +Python: + from ultralytics import YOLO + model = YOLO.new('yolov8n.yaml') + results = model.export(format='onnx') + +CLI: + $ yolo mode=export model=yolov8n.pt format=onnx Inference: $ python detect.py --weights yolov8n.pt # PyTorch yolov8n.torchscript # TorchScript yolov8n.onnx # ONNX Runtime or OpenCV DNN with --dnn - yolov5s_openvino_model # OpenVINO + yolov8n_openvino_model # OpenVINO yolov8n.engine # TensorRT yolov8n.mlmodel # CoreML (macOS-only) - yolov5s_saved_model # TensorFlow SavedModel + yolov8n_saved_model # TensorFlow SavedModel yolov8n.pb # TensorFlow GraphDef yolov8n.tflite # TensorFlow Lite - yolov5s_edgetpu.tflite # TensorFlow Edge TPU - yolov5s_paddle_model # PaddlePaddle + yolov8n_edgetpu.tflite # TensorFlow Edge TPU + yolov8n_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 + $ ln -s ../../yolov5/yolov8n_web_model public/yolov8n_web_model $ npm start - - -from ultralytics import YOLO -model = YOLO().new('yolov8n.yaml') -results = model.export(format='onnx') """ import contextlib import json @@ -59,15 +59,19 @@ import warnings from copy import deepcopy from pathlib import Path +import hydra +import numpy as np import pandas as pd import torch -from torch.utils.mobile_optimizer import optimize_for_mobile from ultralytics.nn.modules import Detect, Segment -from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel -from ultralytics.yolo.utils import LOGGER, ROOT, colorstr, get_default_args -from ultralytics.yolo.utils.checks import check_imgsz, check_requirements, check_version -from ultralytics.yolo.utils.files import file_size, yaml_save +from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, attempt_load_weights +from ultralytics.yolo.configs import get_config +from ultralytics.yolo.data.dataloaders.stream_loaders import LoadImages +from ultralytics.yolo.data.utils import check_dataset +from ultralytics.yolo.utils import DEFAULT_CONFIG, LOGGER, colorstr, get_default_args +from ultralytics.yolo.utils.checks import check_imgsz, check_requirements, check_version, check_yaml +from ultralytics.yolo.utils.files import file_size, increment_path, yaml_save from ultralytics.yolo.utils.ops import Profile from ultralytics.yolo.utils.torch_utils import select_device, smart_inference_mode @@ -110,499 +114,510 @@ def try_export(inner_func): 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 - - -@try_export -def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX:')): - # YOLOv5 ONNX export - check_requirements('onnx>=1.12.0') - import onnx # noqa - - LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...') - f = file.with_suffix('.onnx') - - output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output0'] - if dynamic: - dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640) - if isinstance(model, SegmentationModel): - dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85) - dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160) - elif isinstance(model, DetectionModel): - dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85) - - torch.onnx.export( - model.cpu() if dynamic else model, # --dynamic only compatible with cpu - im.cpu() if dynamic else im, - f, - verbose=False, - opset_version=opset, - do_constant_folding=True, # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False - input_names=['images'], - output_names=output_names, - dynamic_axes=dynamic or None) - - # Checks - model_onnx = onnx.load(f) # load onnx model - onnx.checker.check_model(model_onnx) # check onnx model - - # Metadata - 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) - - # Simplify - 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, prefix=colorstr('OpenVINO:')): - # YOLOv5 OpenVINO export - check_requirements('openvino-dev') # requires openvino-dev: https://pypi.org/project/openvino-dev/ - import openvino.inference_engine as ie # noqa - - LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...') - f = str(file).replace('.pt', f'_openvino_model{os.sep}') - - cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} --data_type {'FP16' if half else 'FP32'}" - subprocess.run(cmd.split(), check=True, env=os.environ) # export - yaml_save(Path(f) / file.with_suffix('.yaml').name, metadata) # add metadata.yaml - return f, None - - -@try_export -def export_paddle(model, im, file, metadata, prefix=colorstr('PaddlePaddle:')): - # YOLOv5 Paddle export - check_requirements(('paddlepaddle', 'x2paddle')) - import x2paddle # noqa - from x2paddle.convert import pytorch2paddle # noqa - - 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]) # export - yaml_save(Path(f) / file.with_suffix('.yaml').name, metadata) # add metadata.yaml - return f, None - - -@try_export -def export_coreml(model, im, file, int8, half, prefix=colorstr('CoreML:')): - # YOLOv5 CoreML export - check_requirements('coremltools') - import coremltools as ct # noqa - - LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...') - f = file.with_suffix('.mlmodel') - - ts = torch.jit.trace(model, im, strict=False) # TorchScript model - 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: # quantization only supported on macOS - ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode) +class Exporter: + + def __init__(self, config=DEFAULT_CONFIG, overrides={}): + self.args = get_config(config, overrides) + project = self.args.project or f"runs/{self.args.task}" + name = self.args.name or f"{self.args.mode}" + self.save_dir = increment_path(Path(project) / name, exist_ok=self.args.exist_ok) + self.save_dir.mkdir(parents=True, exist_ok=True) + self.imgsz = self.args.imgsz + + @smart_inference_mode() + def __call__(self, model=None): + t = time.time() + format = self.args.format.lower() # to lowercase + fmts = tuple(export_formats()['Argument'][1:]) # available export formats + flags = [x == format for x in fmts] + assert sum(flags), f'ERROR: Invalid format={format}, valid formats are {fmts}' + jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags # export booleans + + # Load PyTorch model + self.device = select_device(self.args.device) + if self.args.half: + assert self.device.type != 'cpu' or coreml, '--half only compatible with GPU export, i.e. use --device 0' + assert not self.args.dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic' + + # Checks + if isinstance(self.imgsz, int): + self.imgsz = [self.imgsz] + self.imgsz *= 2 if len(self.imgsz) == 1 else 1 # expand + if self.args.optimize: + assert self.device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu' + + # Input + self.args.batch_size = 1 # TODO: resolve this issue, default 16 not fit for export + gs = int(max(model.stride)) # grid size (max stride) + imgsz = [check_imgsz(x, gs) for x in self.imgsz] # verify img_size are gs-multiples + im = torch.zeros(self.args.batch_size, 3, *imgsz).to(self.device) # image size(1,3,320,192) BCHW iDetection + file = Path(Path(model.yaml['yaml_file']).name) + + # Update model + model = deepcopy(model) + for p in model.parameters(): + p.requires_grad = False + model.eval() + model = model.fuse() + for k, m in model.named_modules(): + if isinstance(m, (Detect, Segment)): + m.dynamic = self.args.dynamic + m.export = True + + y = None + for _ in range(2): + y = model(im) # dry runs + if self.args.half and not coreml: + im, model = im.half(), model.half() # to FP16 + shape = tuple((y[0] if isinstance(y, tuple) else y).shape) # model output shape + LOGGER.info( + f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)") + + # Warnings + warnings.filterwarnings('ignore', category=torch.jit.TracerWarning) # suppress TracerWarning + warnings.filterwarnings('ignore', category=UserWarning) # suppress shape prim::Constant missing ONNX warning + warnings.filterwarnings('ignore', category=DeprecationWarning) # suppress CoreML np.bool deprecation warning + + # Assign + self.im = im + self.model = model + self.file = file + self.metadata = {'stride': int(max(model.stride)), 'names': model.names} # model metadata + + # Exports + f = [''] * len(fmts) # exported filenames + if jit: # TorchScript + f[0], _ = self._export_torchscript() + if engine: # TensorRT required before ONNX + f[1], _ = self._export_engine() + if onnx or xml: # OpenVINO requires ONNX + f[2], _ = self._export_onnx() + if xml: # OpenVINO + f[3], _ = self._export_openvino() + if coreml: # CoreML + f[4], _ = self._export_coreml() + if any((saved_model, pb, tflite, edgetpu, tfjs)): # TensorFlow formats + assert not isinstance(model, ClassificationModel), 'ClassificationModel TF exports not yet supported.' + nms = False + f[5], s_model = self._export_saved_model(nms=nms or self.args.agnostic_nms or tfjs, + agnostic_nms=self.args.agnostic_nms or tfjs) + if pb or tfjs: # pb prerequisite to tfjs + f[6], _ = self._export_pb(s_model,) + if tflite or edgetpu: + f[7], _ = self._export_tflite(s_model, + int8=self.args.int8 or edgetpu, + data=self.args.data, + nms=nms, + agnostic_nms=self.args.agnostic_nms) + if edgetpu: + f[8], _ = self._export_edgetpu() + self._add_tflite_metadata(f[8] or f[7], num_outputs=len(s_model.outputs)) + if tfjs: + f[9], _ = self._export_tfjs() + if paddle: # PaddlePaddle + f[10], _ = self._export_paddle() + + # Finish + f = [str(x) for x in f if x] # filter out '' and None + if any(f): + cls, det, seg = (isinstance(model, x) + for x in (ClassificationModel, DetectionModel, SegmentationModel)) # type + det &= not seg # segmentation models inherit from SegmentationModel(DetectionModel) + s = "-WARNING ⚠️ not yet supported for YOLOv8 exported models" + task = 'detect' if det else 'segment' if seg else 'classify' if cls else '' + LOGGER.info(f'\nExport complete ({time.time() - t:.1f}s)' + f"\nResults saved to {colorstr('bold', file.parent.resolve())}" + f"\nPredict: yolo task={task} mode=predict model={f[-1]} {s}" + f"\nValidate: yolo task={task} mode=val model={f[-1]} {s}" + f"\nVisualize: https://netron.app") + return f # return list of exported files/dirs + + @try_export + def _export_torchscript(self, prefix=colorstr('TorchScript:')): + # YOLOv5 TorchScript model export + LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...') + f = self.file.with_suffix('.torchscript') + + ts = torch.jit.trace(self.model, self.im, strict=False) + d = {"shape": self.im.shape, "stride": int(max(self.model.stride)), "names": self.model.names} + extra_files = {'config.txt': json.dumps(d)} # torch._C.ExtraFilesMap() + if self.args.optimize: # https://pytorch.org/tutorials/recipes/mobile_interpreter.html + LOGGER.info(f'{prefix} optimizing for mobile...') + from torch.utils.mobile_optimizer import optimize_for_mobile + optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files) else: - LOGGER.info(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:')): - # YOLOv5 TensorRT export https://developer.nvidia.com/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': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012 - 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) # opset 12 - model.model[-1].anchor_grid = grid - else: # TensorRT >= 8 - check_version(trt.__version__, '8.0.0', hard=True) # require tensorrt>=8.0.0 - export_onnx(model, im, file, 12, dynamic, simplify) # opset 12 - 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') # TensorRT engine file - 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 - # config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30) # fix TRT 8.4 deprecation notice - - 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() + ts.save(str(f), _extra_files=extra_files) + return f, None + + @try_export + def _export_onnx(self, prefix=colorstr('ONNX:')): + # YOLOv5 ONNX export + check_requirements('onnx>=1.12.0') + import onnx # noqa + + LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...') + f = str(self.file.with_suffix('.onnx')) + + output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output0'] + dynamic = self.args.dynamic + if dynamic: + dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640) + if isinstance(self.model, SegmentationModel): + dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85) + dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160) + elif isinstance(self.model, DetectionModel): + dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85) + + torch.onnx.export( + self.model.cpu() if dynamic else self.model, # --dynamic only compatible with cpu + self.im.cpu() if dynamic else self.im, + f, + verbose=False, + opset_version=self.args.opset, + do_constant_folding=True, # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False + input_names=['images'], + output_names=output_names, + dynamic_axes=dynamic or None) + + # Checks + model_onnx = onnx.load(f) # load onnx model + onnx.checker.check_model(model_onnx) # check onnx model + + # Metadata + d = {'stride': int(max(self.model.stride)), 'names': self.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) + + # Simplify + if self.args.simplify: + try: + cuda = torch.cuda.is_available() + check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1')) + import onnxsim # noqa + + 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(self, prefix=colorstr('OpenVINO:')): + # YOLOv5 OpenVINO export + check_requirements('openvino-dev') # requires openvino-dev: https://pypi.org/project/openvino-dev/ + import openvino.inference_engine as ie # noqa + + LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...') + f = str(self.file).replace(self.file.suffix, f'_openvino_model{os.sep}') + f_onnx = self.file.with_suffix('.onnx') + + cmd = f"mo --input_model {f_onnx} --output_dir {f} --data_type {'FP16' if self.args.half else 'FP32'}" + subprocess.run(cmd.split(), check=True, env=os.environ) # export + yaml_save(Path(f) / self.file.with_suffix('.yaml').name, self.metadata) # add metadata.yaml + return f, None + + @try_export + def _export_paddle(self, prefix=colorstr('PaddlePaddle:')): + # YOLOv5 Paddle export + check_requirements(('paddlepaddle', 'x2paddle')) + import x2paddle # noqa + from x2paddle.convert import pytorch2paddle # noqa + + LOGGER.info(f'\n{prefix} starting export with X2Paddle {x2paddle.__version__}...') + f = str(self.file).replace(self.file.suffix, f'_paddle_model{os.sep}') + + pytorch2paddle(module=self.model, save_dir=f, jit_type='trace', input_examples=[self.im]) # export + yaml_save(Path(f) / self.file.with_suffix('.yaml').name, self.metadata) # add metadata.yaml + return f, None + + @try_export + def _export_coreml(self, prefix=colorstr('CoreML:')): + # YOLOv5 CoreML export + check_requirements('coremltools') + import coremltools as ct # noqa + + LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...') + f = self.file.with_suffix('.mlmodel') + + ts = torch.jit.trace(self.model, self.im, strict=False) # TorchScript model + ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=self.im.shape, scale=1 / 255, bias=[0, 0, 0])]) + bits, mode = (8, 'kmeans_lut') if self.args.int8 else (16, 'linear') if self.args.half else (32, None) + if bits < 32: + if MACOS: # quantization only supported on macOS + ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode) + else: + LOGGER.info(f'{prefix} quantization only supported on macOS, skipping...') + ct_model.save(str(f)) + return f, ct_model + + @try_export + def _export_engine(self, workspace=4, verbose=False, prefix=colorstr('TensorRT:')): + # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt + assert self.im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `device==0`' + try: + import tensorrt as trt # noqa + except ImportError: + if platform.system() == 'Linux': + check_requirements('nvidia-tensorrt', cmds='-U --index-url https://pypi.ngc.nvidia.com') + import tensorrt as trt # noqa + + check_version(trt.__version__, '7.0.0', hard=True) # require tensorrt>=8.0.0 + self._export_onnx() + onnx = self.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 = self.file.with_suffix('.engine') # TensorRT engine file + 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 + # config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30) # fix TRT 8.4 deprecation notice + + 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: - 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:')): - # YOLOv5 TensorFlow SavedModel export - 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 models.tf import TFModel - from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 # noqa - - 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) # BCHW - - tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz) - im = tf.zeros((batch_size, *imgsz, ch)) # BHWC order for TensorFlow - _ = 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) + 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 self.args.dynamic: + shape = self.im.shape + if 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, *shape[1:]), (max(1, shape[0] // 2), *shape[1:]), shape) + config.add_optimization_profile(profile) + + LOGGER.info( + f'{prefix} building FP{16 if builder.platform_has_fast_fp16 and self.args.half else 32} engine as {f}') + if builder.platform_has_fast_fp16 and self.args.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(self, + nms=False, + agnostic_nms=False, + topk_per_class=100, + topk_all=100, + iou_thres=0.45, + conf_thres=0.25, + prefix=colorstr('TensorFlow SavedModel:')): + # YOLOv5 TensorFlow SavedModel export + try: + import tensorflow as tf # noqa + except ImportError: + check_requirements(f"tensorflow{'' if torch.cuda.is_available() else '-macos' if MACOS else '-cpu'}") + import tensorflow as tf # noqa + # from models.tf import TFModel + from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 # noqa + + LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') + f = str(self.file).replace(self.file.suffix, '_saved_model') + batch_size, ch, *imgsz = list(self.im.shape) # BCHW + + tf_models = None # TODO: no TF modules available + tf_model = tf_models.TFModel(cfg=self.model.yaml, model=self.model.cpu(), nc=self.model.nc, imgsz=imgsz) + im = tf.zeros((batch_size, *imgsz, ch)) # BHWC order for TensorFlow + _ = tf_model.predict(im, nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres) + inputs = tf.keras.Input(shape=(*imgsz, ch), batch_size=None if self.args.dynamic else batch_size) + outputs = tf_model.predict(inputs, 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 self.args.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)) # full model + 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 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(self, keras_model, file, prefix=colorstr('TensorFlow GraphDef:')): + # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow + import tensorflow as tf # noqa + from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 # noqa + + LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') + f = file.with_suffix('.pb') + m = tf.function(lambda x: keras_model(x)) # full model - m = m.get_concrete_function(spec) + 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) - 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:')): - # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow - import tensorflow as tf # noqa - from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 # noqa - - LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') - f = file.with_suffix('.pb') - - m = tf.function(lambda x: keras_model(x)) # full model - 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, data, nms, agnostic_nms, prefix=colorstr('TensorFlow Lite:')): - # YOLOv5 TensorFlow Lite export - import tensorflow as tf # noqa - - LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') - batch_size, ch, *imgsz = list(im.shape) # BCHW - 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'], imgsz=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 # or tf.int8 - converter.inference_output_type = tf.uint8 # or tf.int8 - converter.experimental_new_quantizer = 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:')): - # YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/ - 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', 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 # sudo installed on system - 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') # Edge TPU model - f_tfl = str(file).replace('.pt', '-int8.tflite') # TFLite model - - cmd = f"edgetpu_compiler -s -d -k 10 --out_dir {file.parent} {f_tfl}" - subprocess.run(cmd.split(), check=True) - return f, None - - -@try_export -def export_tfjs(file, prefix=colorstr('TensorFlow.js:')): - # YOLOv5 TensorFlow.js export - check_requirements('tensorflowjs') - import tensorflowjs as tfjs # noqa - - LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...') - f = str(file).replace('.pt', '_web_model') # js dir - f_pb = file.with_suffix('.pb') # *.pb path - f_json = f'{f}/model.json' # *.json path - - 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.split()) - - json = Path(f_json).read_text() - with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order - 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): - # Add metadata to *.tflite models per https://www.tensorflow.org/lite/models/convert/metadata - with contextlib.suppress(ImportError): - # check_requirements('tflite_support') - from tflite_support import flatbuffers # noqa - from tflite_support import metadata as _metadata # noqa - from tflite_support import metadata_schema_py_generated as _metadata_fb # noqa - - 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() - - -@smart_inference_mode() -def export_model( - model, # model - file=ROOT / 'yolov8n.pt', - data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' - imgsz=(640, 640), # image (height, width) - batch_size=1, # batch size - device=torch.device('cpu'), # cuda device, i.e. 0 or 0,1,2,3 or cpu - format='onnx', # export format - half=False, # FP16 half-precision export - keras=False, # use Keras - optimize=False, # TorchScript: optimize for mobile - int8=False, # CoreML/TF INT8 quantization - dynamic=False, # ONNX/TF/TensorRT: dynamic axes - simplify=False, # ONNX: simplify model - opset=17, # ONNX: opset version - verbose=False, # TensorRT: verbose log - workspace=4, # TensorRT: workspace size (GB) - nms=False, # TF: add NMS to model - agnostic_nms=False, # TF: add agnostic NMS to model - topk_per_class=100, # TF.js NMS: topk per class to keep - topk_all=100, # TF.js NMS: topk for all classes to keep - iou_thres=0.45, # TF.js NMS: IoU threshold - conf_thres=0.25, # TF.js NMS: confidence threshold -): - t = time.time() - format = format.lower() # to lowercase - fmts = tuple(export_formats()['Argument'][1:]) # available export formats - flags = [x == format for x in fmts] - assert sum(flags), f'ERROR: Invalid format={format}, valid formats are {fmts}' - jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags # export booleans - - # Load PyTorch model - 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 = deepcopy(model).fuse() # load FP32 model - - # Checks - if isinstance(imgsz, int): - imgsz = [imgsz] - imgsz *= 2 if len(imgsz) == 1 else 1 # expand - if optimize: - assert device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu' - - # Input - gs = int(max(model.stride)) # grid size (max stride) - imgsz = [check_imgsz(x, gs) for x in imgsz] # verify img_size are gs-multiples - im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection - - # Update model - model.eval() - for k, m in model.named_modules(): - if isinstance(m, (Detect, Segment)): - m.dynamic = dynamic - m.export = True - - for _ in range(2): - y = model(im) # dry runs - if half and not coreml: - im, model = im.half(), model.half() # to FP16 - shape = tuple((y[0] if isinstance(y, tuple) else y).shape) # model output shape - metadata = {'stride': int(max(model.stride)), 'names': model.names} # model metadata - LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)") - - # Warnings - warnings.filterwarnings('ignore', category=torch.jit.TracerWarning) # suppress TracerWarning - warnings.filterwarnings('ignore', category=UserWarning) # suppress shape prim::Constant type missing ONNX warning - warnings.filterwarnings('ignore', category=DeprecationWarning) # suppress CoreML np.bool deprecation warning - - # Exports - f = [''] * len(fmts) # exported filenames - if jit: # TorchScript - f[0], _ = export_torchscript(model, im, file, optimize) - if engine: # TensorRT required before ONNX - f[1], _ = export_engine(model, im, file, half, dynamic, simplify, workspace, verbose) - if onnx or xml: # OpenVINO requires ONNX - f[2], _ = export_onnx(model, im, file, opset, dynamic, simplify) - if xml: # OpenVINO - f[3], _ = export_openvino(file, metadata, half) - if coreml: # CoreML - f[4], _ = export_coreml(model, im, file, int8, half) - if any((saved_model, pb, tflite, edgetpu, tfjs)): # TensorFlow formats - 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: # pb prerequisite to tfjs - f[6], _ = export_pb(s_model, file) - if tflite or edgetpu: - f[7], _ = export_tflite(s_model, im, file, int8 or edgetpu, 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) - if paddle: # PaddlePaddle - f[10], _ = export_paddle(model, im, file, metadata) - - # Finish - f = [str(x) for x in f if x] # filter out '' and None - if any(f): - cls, det, seg = (isinstance(model, x) for x in (ClassificationModel, DetectionModel, SegmentationModel)) # type - det &= not seg # segmentation models inherit from SegmentationModel(DetectionModel) - dir = Path('segment' if seg else 'classify' if cls else '') - h = '--half' if half else '' # --half FP16 inference arg - 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 / '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 # return list of exported files/dirs + 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(self, keras_model, int8, data, nms, agnostic_nms, prefix=colorstr('TensorFlow Lite:')): + # YOLOv5 TensorFlow Lite export + import tensorflow as tf # noqa + + LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') + batch_size, ch, *imgsz = list(self.im.shape) # BCHW + f = str(self.file).replace(self.file.suffix, '-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: + + def representative_dataset_gen(dataset, n_images=100): + # Dataset generator for use with converter.representative_dataset, returns a generator of np arrays + for n, (path, img, im0s, vid_cap, string) in enumerate(dataset): + im = np.transpose(img, [1, 2, 0]) + im = np.expand_dims(im, axis=0).astype(np.float32) + im /= 255 + yield [im] + if n >= n_images: + break + + dataset = LoadImages(check_dataset(check_yaml(data))['train'], imgsz=imgsz, auto=False) + converter.representative_dataset = lambda: representative_dataset_gen(dataset, n_images=100) + converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] + converter.target_spec.supported_types = [] + converter.inference_input_type = tf.uint8 # or tf.int8 + converter.inference_output_type = tf.uint8 # or tf.int8 + converter.experimental_new_quantizer = True + f = str(self.file).replace(self.file.suffix, '-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(self, prefix=colorstr('Edge TPU:')): + # YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/ + 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', 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 # sudo installed on system + 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(self.file).replace(self.file.suffix, '-int8_edgetpu.tflite') # Edge TPU model + f_tfl = str(self.file).replace(self.file.suffix, '-int8.tflite') # TFLite model + + cmd = f"edgetpu_compiler -s -d -k 10 --out_dir {self.file.parent} {f_tfl}" + subprocess.run(cmd.split(), check=True) + return f, None + + @try_export + def _export_tfjs(self, prefix=colorstr('TensorFlow.js:')): + # YOLOv5 TensorFlow.js export + check_requirements('tensorflowjs') + import tensorflowjs as tfjs # noqa + + LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...') + f = str(self.file).replace(self.file.suffix, '_web_model') # js dir + f_pb = self.file.with_suffix('.pb') # *.pb path + f_json = Path(f) / 'model.json' # *.json path + + 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.split()) + + with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order + 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"}}}', f_json.read_text()) + j.write(subst) + return f, None + + def _add_tflite_metadata(self, file, num_outputs): + # Add metadata to *.tflite models per https://www.tensorflow.org/lite/models/convert/metadata + with contextlib.suppress(ImportError): + # check_requirements('tflite_support') + from tflite_support import flatbuffers # noqa + from tflite_support import metadata as _metadata # noqa + from tflite_support import metadata_schema_py_generated as _metadata_fb # noqa + + tmp_file = Path('/tmp/meta.txt') + with open(tmp_file, 'w') as meta_f: + meta_f.write(str(self.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() + + +@hydra.main(version_base=None, config_path=str(DEFAULT_CONFIG.parent), config_name=DEFAULT_CONFIG.name) +def export(cfg): + cfg.model = cfg.model or "yolov8n.yaml" + cfg.format = cfg.format or "torchscript" + exporter = Exporter(cfg) + + model = None + if isinstance(cfg.model, (str, Path)): + if Path(cfg.model).suffix == '.yaml': + model = DetectionModel(cfg.model) + elif Path(cfg.model).suffix == '.pt': + model = attempt_load_weights(cfg.model) + else: + TypeError(f'Unsupported model type {cfg.model}') + exporter(model=model) + + +if __name__ == "__main__": + """ + CLI: + yolo mode=export model=yolov8n.yaml format=onnx + """ + export() diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py index 05b2d5b..025fef6 100644 --- a/ultralytics/yolo/engine/model.py +++ b/ultralytics/yolo/engine/model.py @@ -5,7 +5,7 @@ import torch from ultralytics import yolo # noqa required for python usage from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, attempt_load_weights from ultralytics.yolo.configs import get_config -from ultralytics.yolo.engine.exporter import export_model +from ultralytics.yolo.engine.exporter import Exporter from ultralytics.yolo.utils import DEFAULT_CONFIG, HELP_MSG, LOGGER from ultralytics.yolo.utils.checks import check_yaml from ultralytics.yolo.utils.files import yaml_load @@ -164,7 +164,7 @@ class YOLO: validator(model=self.model) @smart_inference_mode() - def export(self, format='', save_dir='', **kwargs): + def export(self, **kwargs): """ Export model. @@ -177,36 +177,9 @@ class YOLO: overrides.update(kwargs) args = get_config(config=DEFAULT_CONFIG, overrides=overrides) args.task = self.task - args.format = format - - file = self.ckpt or Path(Path(self.cfg).name) - if save_dir: - file = Path(save_dir) / file.name - file.parent.mkdir(parents=True, exist_ok=True) - - export_model( - model=self.model, - file=file, - data=args.data, # 'dataset.yaml path' - imgsz=args.imgsz or (640, 640), # image (height, width) - batch_size=1, # batch size - device=args.device, # cuda device, i.e. 0 or 0,1,2,3 or cpu - format=args.format, # include formats - half=args.half or False, # FP16 half-precision export - keras=args.keras or False, # use Keras - optimize=args.optimize or False, # TorchScript: optimize for mobile - int8=args.int8 or False, # CoreML/TF INT8 quantization - dynamic=args.dynamic or False, # ONNX/TF/TensorRT: dynamic axes - opset=args.opset or 17, # ONNX: opset version - verbose=False, # TensorRT: verbose log - workspace=args.workspace or 4, # TensorRT: workspace size (GB) - nms=False, # TF: add NMS to model - agnostic_nms=False, # TF: add agnostic NMS to model - topk_per_class=100, # TF.js NMS: topk per class to keep - topk_all=100, # TF.js NMS: topk for all classes to keep - iou_thres=0.45, # TF.js NMS: IoU threshold - conf_thres=0.25, # TF.js NMS: confidence threshold - ) + + exporter = Exporter(overrides=overrides) + exporter(model=self.model) def train(self, **kwargs): """ diff --git a/ultralytics/yolo/engine/predictor.py b/ultralytics/yolo/engine/predictor.py index 641d05f..62f44be 100644 --- a/ultralytics/yolo/engine/predictor.py +++ b/ultralytics/yolo/engine/predictor.py @@ -16,14 +16,14 @@ Usage - formats: $ yolo task=... mode=predict --weights yolov8n.pt # PyTorch yolov8n.torchscript # TorchScript yolov8n.onnx # ONNX Runtime or OpenCV DNN with --dnn - yolov5s_openvino_model # OpenVINO + yolov8n_openvino_model # OpenVINO yolov8n.engine # TensorRT yolov8n.mlmodel # CoreML (macOS-only) - yolov5s_saved_model # TensorFlow SavedModel + yolov8n_saved_model # TensorFlow SavedModel yolov8n.pb # TensorFlow GraphDef yolov8n.tflite # TensorFlow Lite - yolov5s_edgetpu.tflite # TensorFlow Edge TPU - yolov5s_paddle_model # PaddlePaddle + yolov8n_edgetpu.tflite # TensorFlow Edge TPU + yolov8n_paddle_model # PaddlePaddle """ import platform from pathlib import Path diff --git a/ultralytics/yolo/utils/__init__.py b/ultralytics/yolo/utils/__init__.py index ecbdfe6..21b4e1f 100644 --- a/ultralytics/yolo/utils/__init__.py +++ b/ultralytics/yolo/utils/__init__.py @@ -25,14 +25,12 @@ TQDM_BAR_FORMAT = '{l_bar}{bar:10}{r_bar}' # tqdm bar format LOGGING_NAME = 'yolov5' HELP_MSG = \ """ - Please refer to below Usage examples for help running YOLOv8 - For help visit Ultralytics Community at https://community.ultralytics.com/ - Submit bug reports to https//github.com/ultralytics/ultralytics + Please refer to below Usage examples for help running YOLOv8: Install: pip install ultralytics - Python usage: + Python SDK: from ultralytics import YOLO model = YOLO.new('yolov8n.yaml') # create a new model from scratch @@ -42,12 +40,15 @@ HELP_MSG = \ results = model.predict(source='bus.jpg') success = model.export(format='onnx') - CLI usage: - yolo task=detect mode=train model=yolov8n.yaml ... - classify predict yolov8n-cls.yaml - segment val yolov8n-seg.yaml + CLI: + yolo task=detect mode=train model=yolov8n.yaml args... + classify predict yolov8n-cls.yaml args... + segment val yolov8n-seg.yaml args... + export yolov8n.pt format=onnx args... - For all arguments see https://github.com/ultralytics/ultralytics/blob/main/ultralytics/yolo/utils/configs/default.yaml + Docs: https://docs.ultralytics.com + Community: https://community.ultralytics.com + GitHub: https://github.com/ultralytics/ultralytics """ # Settings @@ -56,7 +57,6 @@ HELP_MSG = \ pd.options.display.max_columns = 10 cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader) os.environ['NUMEXPR_MAX_THREADS'] = str(NUM_THREADS) # NumExpr max threads -os.environ['OMP_NUM_THREADS'] = '1' if platform.system() == 'darwin' else str(NUM_THREADS) # OpenMP (PyTorch and SciPy) def is_colab(): diff --git a/ultralytics/yolo/utils/callbacks/clearml.py b/ultralytics/yolo/utils/callbacks/clearml.py index 8c01cbd..225a25f 100644 --- a/ultralytics/yolo/utils/callbacks/clearml.py +++ b/ultralytics/yolo/utils/callbacks/clearml.py @@ -36,8 +36,8 @@ def on_val_end(trainer): if trainer.epoch == 0: model_info = { "Parameters": get_num_params(trainer.model), - "GFLOPs": round(get_flops(trainer.model), 1), - "Inference speed (ms/img)": round(trainer.validator.speed[1], 1)} + "GFLOPs": round(get_flops(trainer.model), 3), + "Inference speed (ms/img)": round(trainer.validator.speed[1], 3)} Task.current_task().connect(model_info, name='Model') diff --git a/ultralytics/yolo/utils/callbacks/wb.py b/ultralytics/yolo/utils/callbacks/wb.py index d0287d1..eed8ead 100644 --- a/ultralytics/yolo/utils/callbacks/wb.py +++ b/ultralytics/yolo/utils/callbacks/wb.py @@ -19,8 +19,8 @@ def on_val_end(trainer): if trainer.epoch == 0: model_info = { "model/parameters": get_num_params(trainer.model), - "model/GFLOPs": round(get_flops(trainer.model), 1), - "model/speed(ms)": round(trainer.validator.speed[1], 1)} + "model/GFLOPs": round(get_flops(trainer.model), 3), + "model/speed(ms)": round(trainer.validator.speed[1], 3)} wandb.run.log(model_info, step=trainer.epoch + 1)