xdobro23
/
YOLOv8-16bit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Create Exporter() Class (#117)

Browse Source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
single_channel
Glenn Jocher 2 years ago committed by GitHub
parent a9dc1637c2
commit 076d73cfaa
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 563 additions and 572 deletions
Show Stats Download Patch File Download Diff File

7
README.md
Unescape View File

@ -16,9 +16,10 @@ pip install -e .
### 1. CLI ### 1. CLI
To simply use the latest Ultralytics YOLO models To simply use the latest Ultralytics YOLO models
```bash ```bash
yolo task=detect mode=train model=yolov8n.yaml ... yolo task=detect mode=train model=yolov8n.yaml args=...
classify predict yolov8n-cls.yaml classify predict yolov8n-cls.yaml args=...
segment val yolov8n-seg.yaml segment val yolov8n-seg.yaml args=...
export yolov8n.pt format=onnx
``` ```
### 2. Python SDK ### 2. Python SDK
To use pythonic interface of Ultralytics YOLO model To use pythonic interface of Ultralytics YOLO model

3
ultralytics/nn/tasks.py
Unescape View File

@ -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, Concat, Conv, ConvTranspose, Detect, DWConv, DWConvTranspose2d, Ensemble, Focus,
GhostBottleneck, GhostConv, Segment) GhostBottleneck, GhostConv, Segment)
from ultralytics.yolo.utils import LOGGER, colorstr 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.files import yaml_load
from ultralytics.yolo.utils.torch_utils import (fuse_conv_and_bn, initialize_weights, intersect_state_dicts, from ultralytics.yolo.utils.torch_utils import (fuse_conv_and_bn, initialize_weights, intersect_state_dicts,
make_divisible, model_info, scale_img, time_sync) make_divisible, model_info, scale_img, time_sync)
@ -80,7 +81,7 @@ class DetectionModel(BaseModel):
# YOLOv5 detection model # YOLOv5 detection model
def __init__(self, cfg='yolov8n.yaml', ch=3, nc=None, verbose=True): # model, input channels, number of classes def __init__(self, cfg='yolov8n.yaml', ch=3, nc=None, verbose=True): # model, input channels, number of classes
super().__init__() 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 # Define model
ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels

4
ultralytics/yolo/cli.py
Unescape View File

@ -31,7 +31,7 @@ def cli(cfg):
elif task == "classify": elif task == "classify":
module = yolo.v8.classify module = yolo.v8.classify
elif task == "export": elif task == "export":
func = yolo.trainer.exporter.export_model func = yolo.engine.exporter.export
else: else:
raise SyntaxError("task not recognized. Choices are `'detect', 'segment', 'classify'`") raise SyntaxError("task not recognized. Choices are `'detect', 'segment', 'classify'`")
@ -42,7 +42,7 @@ def cli(cfg):
elif mode == "predict": elif mode == "predict":
func = module.predict func = module.predict
elif mode == "export": elif mode == "export":
func = yolo.trainer.exporter.export_model func = yolo.engine.exporter.export
else: else:
raise SyntaxError("mode not recognized. Choices are `'train', 'val', 'predict', 'export'`") raise SyntaxError("mode not recognized. Choices are `'train', 'val', 'predict', 'export'`")
func(cfg) func(cfg)

3
ultralytics/yolo/configs/default.yaml
Unescape View File

@ -29,12 +29,12 @@ image_weights: False # use weighted image selection for training
rect: False # support rectangular training rect: False # support rectangular training
cos_lr: False # use cosine LR scheduler cos_lr: False # use cosine LR scheduler
close_mosaic: 10 # disable mosaic for final 10 epochs close_mosaic: 10 # disable mosaic for final 10 epochs
resume: False
# Segmentation # Segmentation
overlap_mask: True # masks overlap overlap_mask: True # masks overlap
mask_ratio: 4 # mask downsample ratio mask_ratio: 4 # mask downsample ratio
# Classification # Classification
dropout: False # use dropout dropout: False # use dropout
resume: False
# Val/Test settings ---------------------------------------------------------------------------------------------------- # Val/Test settings ----------------------------------------------------------------------------------------------------
@ -65,6 +65,7 @@ agnostic_nms: False # class-agnostic NMS
retina_masks: False retina_masks: False
# Export settings ------------------------------------------------------------------------------------------------------ # Export settings ------------------------------------------------------------------------------------------------------
format: torchscript
keras: False # use Keras keras: False # use Keras
optimize: False # TorchScript: optimize for mobile optimize: False # TorchScript: optimize for mobile
int8: False # CoreML/TF INT8 quantization int8: False # CoreML/TF INT8 quantization

531
ultralytics/yolo/engine/exporter.py
Unescape View File

@ -2,51 +2,51 @@
""" """
Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit 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 PyTorch | - | yolov8n.pt
TorchScript | `torchscript` | yolov8n.torchscript TorchScript | `torchscript` | yolov8n.torchscript
ONNX | `onnx` | yolov8n.onnx ONNX | `onnx` | yolov8n.onnx
OpenVINO | `openvino` | yolov5s_openvino_model/ OpenVINO | `openvino` | yolov8n_openvino_model/
TensorRT | `engine` | yolov8n.engine TensorRT | `engine` | yolov8n.engine
CoreML | `coreml` | yolov8n.mlmodel CoreML | `coreml` | yolov8n.mlmodel
TensorFlow SavedModel | `saved_model` | yolov5s_saved_model/ TensorFlow SavedModel | `saved_model` | yolov8n_saved_model/
TensorFlow GraphDef | `pb` | yolov8n.pb TensorFlow GraphDef | `pb` | yolov8n.pb
TensorFlow Lite | `tflite` | yolov8n.tflite TensorFlow Lite | `tflite` | yolov8n.tflite
TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite TensorFlow Edge TPU | `edgetpu` | yolov8n_edgetpu.tflite
TensorFlow.js | `tfjs` | yolov5s_web_model/ TensorFlow.js | `tfjs` | yolov8n_web_model/
PaddlePaddle | `paddle` | yolov5s_paddle_model/ PaddlePaddle | `paddle` | yolov8n_paddle_model/
Requirements: 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 openvino-dev tensorflow-cpu # CPU
$ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU
Usage: Python:
$ python export.py --weights yolov8n.pt --include torchscript onnx openvino engine coreml tflite ... from ultralytics import YOLO
model = YOLO.new('yolov8n.yaml')
results = model.export(format='onnx')
CLI:
$ yolo mode=export model=yolov8n.pt format=onnx
Inference: Inference:
$ python detect.py --weights yolov8n.pt # PyTorch $ python detect.py --weights yolov8n.pt # PyTorch
yolov8n.torchscript # TorchScript yolov8n.torchscript # TorchScript
yolov8n.onnx # ONNX Runtime or OpenCV DNN with --dnn yolov8n.onnx # ONNX Runtime or OpenCV DNN with --dnn
yolov5s_openvino_model # OpenVINO yolov8n_openvino_model # OpenVINO
yolov8n.engine # TensorRT yolov8n.engine # TensorRT
yolov8n.mlmodel # CoreML (macOS-only) yolov8n.mlmodel # CoreML (macOS-only)
yolov5s_saved_model # TensorFlow SavedModel yolov8n_saved_model # TensorFlow SavedModel
yolov8n.pb # TensorFlow GraphDef yolov8n.pb # TensorFlow GraphDef
yolov8n.tflite # TensorFlow Lite yolov8n.tflite # TensorFlow Lite
yolov5s_edgetpu.tflite # TensorFlow Edge TPU yolov8n_edgetpu.tflite # TensorFlow Edge TPU
yolov5s_paddle_model # PaddlePaddle yolov8n_paddle_model # PaddlePaddle
TensorFlow.js: TensorFlow.js:
$ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
$ npm install $ npm install
$ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model $ ln -s ../../yolov5/yolov8n_web_model public/yolov8n_web_model
$ npm start $ npm start
from ultralytics import YOLO
model = YOLO().new('yolov8n.yaml')
results = model.export(format='onnx')
""" """
import contextlib import contextlib
import json import json
@ -59,15 +59,19 @@ import warnings
from copy import deepcopy from copy import deepcopy
from pathlib import Path from pathlib import Path
import hydra
import numpy as np
import pandas as pd import pandas as pd
import torch import torch
from torch.utils.mobile_optimizer import optimize_for_mobile
from ultralytics.nn.modules import Detect, Segment from ultralytics.nn.modules import Detect, Segment
from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, attempt_load_weights
from ultralytics.yolo.utils import LOGGER, ROOT, colorstr, get_default_args from ultralytics.yolo.configs import get_config
from ultralytics.yolo.utils.checks import check_imgsz, check_requirements, check_version from ultralytics.yolo.data.dataloaders.stream_loaders import LoadImages
from ultralytics.yolo.utils.files import file_size, yaml_save 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.ops import Profile
from ultralytics.yolo.utils.torch_utils import select_device, smart_inference_mode from ultralytics.yolo.utils.torch_utils import select_device, smart_inference_mode
@ -110,46 +114,166 @@ def try_export(inner_func):
return outer_func return outer_func
@try_export class Exporter:
def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
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 # YOLOv5 TorchScript model export
LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...') LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...')
f = file.with_suffix('.torchscript') f = self.file.with_suffix('.torchscript')
ts = torch.jit.trace(model, im, strict=False) ts = torch.jit.trace(self.model, self.im, strict=False)
d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names} 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() extra_files = {'config.txt': json.dumps(d)} # torch._C.ExtraFilesMap()
if optimize: # https://pytorch.org/tutorials/recipes/mobile_interpreter.html 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) optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)
else: else:
ts.save(str(f), _extra_files=extra_files) ts.save(str(f), _extra_files=extra_files)
return f, None return f, None
@try_export
@try_export def _export_onnx(self, prefix=colorstr('ONNX:')):
def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX:')):
# YOLOv5 ONNX export # YOLOv5 ONNX export
check_requirements('onnx>=1.12.0') check_requirements('onnx>=1.12.0')
import onnx # noqa import onnx # noqa
LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...') LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
f = file.with_suffix('.onnx') f = str(self.file.with_suffix('.onnx'))
output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output0'] output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output0']
dynamic = self.args.dynamic
if dynamic: if dynamic:
dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640) dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640)
if isinstance(model, SegmentationModel): if isinstance(self.model, SegmentationModel):
dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85) 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) dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
elif isinstance(model, DetectionModel): elif isinstance(self.model, DetectionModel):
dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85) dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
torch.onnx.export( torch.onnx.export(
model.cpu() if dynamic else model, # --dynamic only compatible with cpu self.model.cpu() if dynamic else self.model, # --dynamic only compatible with cpu
im.cpu() if dynamic else im, self.im.cpu() if dynamic else self.im,
f, f,
verbose=False, verbose=False,
opset_version=opset, opset_version=self.args.opset,
do_constant_folding=True, # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False do_constant_folding=True, # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
input_names=['images'], input_names=['images'],
output_names=output_names, output_names=output_names,
@ -160,18 +284,18 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX
onnx.checker.check_model(model_onnx) # check onnx model onnx.checker.check_model(model_onnx) # check onnx model
# Metadata # Metadata
d = {'stride': int(max(model.stride)), 'names': model.names} d = {'stride': int(max(self.model.stride)), 'names': self.model.names}
for k, v in d.items(): for k, v in d.items():
meta = model_onnx.metadata_props.add() meta = model_onnx.metadata_props.add()
meta.key, meta.value = k, str(v) meta.key, meta.value = k, str(v)
onnx.save(model_onnx, f) onnx.save(model_onnx, f)
# Simplify # Simplify
if simplify: if self.args.simplify:
try: try:
cuda = torch.cuda.is_available() cuda = torch.cuda.is_available()
check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1')) check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1'))
import onnxsim import onnxsim # noqa
LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...') LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
model_onnx, check = onnxsim.simplify(model_onnx) model_onnx, check = onnxsim.simplify(model_onnx)
@ -181,82 +305,73 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX
LOGGER.info(f'{prefix} simplifier failure: {e}') LOGGER.info(f'{prefix} simplifier failure: {e}')
return f, model_onnx return f, model_onnx
@try_export
@try_export def _export_openvino(self, prefix=colorstr('OpenVINO:')):
def export_openvino(file, metadata, half, prefix=colorstr('OpenVINO:')):
# YOLOv5 OpenVINO export # YOLOv5 OpenVINO export
check_requirements('openvino-dev') # requires openvino-dev: https://pypi.org/project/openvino-dev/ check_requirements('openvino-dev') # requires openvino-dev: https://pypi.org/project/openvino-dev/
import openvino.inference_engine as ie # noqa import openvino.inference_engine as ie # noqa
LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...') LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...')
f = str(file).replace('.pt', f'_openvino_model{os.sep}') 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 {file.with_suffix('.onnx')} --output_dir {f} --data_type {'FP16' if half else 'FP32'}" 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 subprocess.run(cmd.split(), check=True, env=os.environ) # export
yaml_save(Path(f) / file.with_suffix('.yaml').name, metadata) # add metadata.yaml yaml_save(Path(f) / self.file.with_suffix('.yaml').name, self.metadata) # add metadata.yaml
return f, None return f, None
@try_export
@try_export def _export_paddle(self, prefix=colorstr('PaddlePaddle:')):
def export_paddle(model, im, file, metadata, prefix=colorstr('PaddlePaddle:')):
# YOLOv5 Paddle export # YOLOv5 Paddle export
check_requirements(('paddlepaddle', 'x2paddle')) check_requirements(('paddlepaddle', 'x2paddle'))
import x2paddle # noqa import x2paddle # noqa
from x2paddle.convert import pytorch2paddle # noqa from x2paddle.convert import pytorch2paddle # noqa
LOGGER.info(f'\n{prefix} starting export with X2Paddle {x2paddle.__version__}...') LOGGER.info(f'\n{prefix} starting export with X2Paddle {x2paddle.__version__}...')
f = str(file).replace('.pt', f'_paddle_model{os.sep}') f = str(self.file).replace(self.file.suffix, f'_paddle_model{os.sep}')
pytorch2paddle(module=model, save_dir=f, jit_type='trace', input_examples=[im]) # export pytorch2paddle(module=self.model, save_dir=f, jit_type='trace', input_examples=[self.im]) # export
yaml_save(Path(f) / file.with_suffix('.yaml').name, metadata) # add metadata.yaml yaml_save(Path(f) / self.file.with_suffix('.yaml').name, self.metadata) # add metadata.yaml
return f, None return f, None
@try_export
@try_export def _export_coreml(self, prefix=colorstr('CoreML:')):
def export_coreml(model, im, file, int8, half, prefix=colorstr('CoreML:')):
# YOLOv5 CoreML export # YOLOv5 CoreML export
check_requirements('coremltools') check_requirements('coremltools')
import coremltools as ct # noqa import coremltools as ct # noqa
LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...') LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...')
f = file.with_suffix('.mlmodel') f = self.file.with_suffix('.mlmodel')
ts = torch.jit.trace(model, im, strict=False) # TorchScript model ts = torch.jit.trace(self.model, self.im, strict=False) # TorchScript model
ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])]) 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 int8 else (16, 'linear') if half else (32, None) bits, mode = (8, 'kmeans_lut') if self.args.int8 else (16, 'linear') if self.args.half else (32, None)
if bits < 32: if bits < 32:
if MACOS: # quantization only supported on macOS if MACOS: # quantization only supported on macOS
ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode) ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)
else: else:
LOGGER.info(f'{prefix} quantization only supported on macOS, skipping...') LOGGER.info(f'{prefix} quantization only supported on macOS, skipping...')
ct_model.save(f) ct_model.save(str(f))
return f, ct_model return f, ct_model
@try_export
@try_export def _export_engine(self, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):
def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):
# YOLOv5 TensorRT export https://developer.nvidia.com/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`' assert self.im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `device==0`'
try: try:
import tensorrt as trt import tensorrt as trt # noqa
except Exception: except ImportError:
if platform.system() == 'Linux': if platform.system() == 'Linux':
check_requirements('nvidia-tensorrt', cmds='-U --index-url https://pypi.ngc.nvidia.com') check_requirements('nvidia-tensorrt', cmds='-U --index-url https://pypi.ngc.nvidia.com')
import tensorrt as trt import tensorrt as trt # noqa
if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012 check_version(trt.__version__, '7.0.0', hard=True) # require tensorrt>=8.0.0
grid = model.model[-1].anchor_grid self._export_onnx()
model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid] onnx = self.file.with_suffix('.onnx')
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__}...') LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
assert onnx.exists(), f'failed to export ONNX file: {onnx}' assert onnx.exists(), f'failed to export ONNX file: {onnx}'
f = file.with_suffix('.engine') # TensorRT engine file f = self.file.with_suffix('.engine') # TensorRT engine file
logger = trt.Logger(trt.Logger.INFO) logger = trt.Logger(trt.Logger.INFO)
if verbose: if verbose:
logger.min_severity = trt.Logger.Severity.VERBOSE logger.min_severity = trt.Logger.Severity.VERBOSE
@ -279,57 +394,55 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose
for out in outputs: for out in outputs:
LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}') LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}')
if dynamic: if self.args.dynamic:
if im.shape[0] <= 1: shape = self.im.shape
if shape[0] <= 1:
LOGGER.warning(f"{prefix} WARNING ⚠️ --dynamic model requires maximum --batch-size argument") LOGGER.warning(f"{prefix} WARNING ⚠️ --dynamic model requires maximum --batch-size argument")
profile = builder.create_optimization_profile() profile = builder.create_optimization_profile()
for inp in inputs: for inp in inputs:
profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape) profile.set_shape(inp.name, (1, *shape[1:]), (max(1, shape[0] // 2), *shape[1:]), shape)
config.add_optimization_profile(profile) 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}') LOGGER.info(
if builder.platform_has_fast_fp16 and half: 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) config.set_flag(trt.BuilderFlag.FP16)
with builder.build_engine(network, config) as engine, open(f, 'wb') as t: with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
t.write(engine.serialize()) t.write(engine.serialize())
return f, None return f, None
@try_export
@try_export def _export_saved_model(self,
def export_saved_model(model, nms=False,
im,
file,
dynamic,
tf_nms=False,
agnostic_nms=False, agnostic_nms=False,
topk_per_class=100, topk_per_class=100,
topk_all=100, topk_all=100,
iou_thres=0.45, iou_thres=0.45,
conf_thres=0.25, conf_thres=0.25,
keras=False,
prefix=colorstr('TensorFlow SavedModel:')): prefix=colorstr('TensorFlow SavedModel:')):
# YOLOv5 TensorFlow SavedModel export # YOLOv5 TensorFlow SavedModel export
try: try:
import tensorflow as tf import tensorflow as tf # noqa
except Exception: except ImportError:
check_requirements(f"tensorflow{'' if torch.cuda.is_available() else '-macos' if MACOS else '-cpu'}") check_requirements(f"tensorflow{'' if torch.cuda.is_available() else '-macos' if MACOS else '-cpu'}")
import tensorflow as tf import tensorflow as tf # noqa
from models.tf import TFModel # from models.tf import TFModel
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 # 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__}...') LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
f = str(file).replace('.pt', '_saved_model') f = str(self.file).replace(self.file.suffix, '_saved_model')
batch_size, ch, *imgsz = list(im.shape) # BCHW batch_size, ch, *imgsz = list(self.im.shape) # BCHW
tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz) 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 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) _ = 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 dynamic else batch_size) inputs = tf.keras.Input(shape=(*imgsz, ch), batch_size=None if self.args.dynamic else batch_size)
outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres) 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 = tf.keras.Model(inputs=inputs, outputs=outputs)
keras_model.trainable = False keras_model.trainable = False
keras_model.summary() keras_model.summary()
if keras: if self.args.keras:
keras_model.save(f, save_format='tf') keras_model.save(f, save_format='tf')
else: else:
spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype) spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)
@ -337,17 +450,16 @@ def export_saved_model(model,
m = m.get_concrete_function(spec) m = m.get_concrete_function(spec)
frozen_func = convert_variables_to_constants_v2(m) frozen_func = convert_variables_to_constants_v2(m)
tfm = tf.Module() tfm = tf.Module()
tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if tf_nms else frozen_func(x), [spec]) tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if nms else frozen_func(x), [spec])
tfm.__call__(im) tfm.__call__(im)
tf.saved_model.save(tfm, tf.saved_model.save(tfm,
f, f,
options=tf.saved_model.SaveOptions(experimental_custom_gradients=False) if check_version( options=tf.saved_model.SaveOptions(experimental_custom_gradients=False)
tf.__version__, '2.6') else tf.saved_model.SaveOptions()) if check_version(tf.__version__, '2.6') else tf.saved_model.SaveOptions())
return f, keras_model return f, keras_model
@try_export
@try_export def _export_pb(self, keras_model, file, prefix=colorstr('TensorFlow GraphDef:')):
def export_pb(keras_model, file, prefix=colorstr('TensorFlow GraphDef:')):
# YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow
import tensorflow as tf # noqa import tensorflow as tf # noqa
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 # noqa from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 # noqa
@ -362,30 +474,39 @@ def export_pb(keras_model, file, prefix=colorstr('TensorFlow GraphDef:')):
tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False) tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
return f, None return f, None
@try_export
@try_export def _export_tflite(self, keras_model, int8, data, nms, agnostic_nms, prefix=colorstr('TensorFlow Lite:')):
def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=colorstr('TensorFlow Lite:')):
# YOLOv5 TensorFlow Lite export # YOLOv5 TensorFlow Lite export
import tensorflow as tf # noqa import tensorflow as tf # noqa
LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
batch_size, ch, *imgsz = list(im.shape) # BCHW batch_size, ch, *imgsz = list(self.im.shape) # BCHW
f = str(file).replace('.pt', '-fp16.tflite') f = str(self.file).replace(self.file.suffix, '-fp16.tflite')
converter = tf.lite.TFLiteConverter.from_keras_model(keras_model) converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS] converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
converter.target_spec.supported_types = [tf.float16] converter.target_spec.supported_types = [tf.float16]
converter.optimizations = [tf.lite.Optimize.DEFAULT] converter.optimizations = [tf.lite.Optimize.DEFAULT]
if int8: if int8:
# from models.tf import representative_dataset_gen
# dataset = LoadImages(check_dataset(check_yaml(data))['train'], imgsz=imgsz, auto=False) def representative_dataset_gen(dataset, n_images=100):
# converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=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_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.target_spec.supported_types = [] converter.target_spec.supported_types = []
converter.inference_input_type = tf.uint8 # or tf.int8 converter.inference_input_type = tf.uint8 # or tf.int8
converter.inference_output_type = tf.uint8 # or tf.int8 converter.inference_output_type = tf.uint8 # or tf.int8
converter.experimental_new_quantizer = True converter.experimental_new_quantizer = True
f = str(file).replace('.pt', '-int8.tflite') f = str(self.file).replace(self.file.suffix, '-int8.tflite')
if nms or agnostic_nms: if nms or agnostic_nms:
converter.target_spec.supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS) converter.target_spec.supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS)
@ -393,9 +514,8 @@ def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=c
open(f, "wb").write(tflite_model) open(f, "wb").write(tflite_model)
return f, None return f, None
@try_export
@try_export def _export_edgetpu(self, prefix=colorstr('Edge TPU:')):
def export_edgetpu(file, prefix=colorstr('Edge TPU:')):
# YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/ # YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/
cmd = 'edgetpu_compiler --version' cmd = 'edgetpu_compiler --version'
help_url = 'https://coral.ai/docs/edgetpu/compiler/' help_url = 'https://coral.ai/docs/edgetpu/compiler/'
@ -411,30 +531,28 @@ def export_edgetpu(file, prefix=colorstr('Edge TPU:')):
ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1] 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}...') LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...')
f = str(file).replace('.pt', '-int8_edgetpu.tflite') # Edge TPU model f = str(self.file).replace(self.file.suffix, '-int8_edgetpu.tflite') # Edge TPU model
f_tfl = str(file).replace('.pt', '-int8.tflite') # TFLite model f_tfl = str(self.file).replace(self.file.suffix, '-int8.tflite') # TFLite model
cmd = f"edgetpu_compiler -s -d -k 10 --out_dir {file.parent} {f_tfl}" cmd = f"edgetpu_compiler -s -d -k 10 --out_dir {self.file.parent} {f_tfl}"
subprocess.run(cmd.split(), check=True) subprocess.run(cmd.split(), check=True)
return f, None return f, None
@try_export
@try_export def _export_tfjs(self, prefix=colorstr('TensorFlow.js:')):
def export_tfjs(file, prefix=colorstr('TensorFlow.js:')):
# YOLOv5 TensorFlow.js export # YOLOv5 TensorFlow.js export
check_requirements('tensorflowjs') check_requirements('tensorflowjs')
import tensorflowjs as tfjs # noqa import tensorflowjs as tfjs # noqa
LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...') LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...')
f = str(file).replace('.pt', '_web_model') # js dir f = str(self.file).replace(self.file.suffix, '_web_model') # js dir
f_pb = file.with_suffix('.pb') # *.pb path f_pb = self.file.with_suffix('.pb') # *.pb path
f_json = f'{f}/model.json' # *.json path f_json = Path(f) / 'model.json' # *.json path
cmd = f'tensorflowjs_converter --input_format=tf_frozen_model ' \ cmd = f'tensorflowjs_converter --input_format=tf_frozen_model ' \
f'--output_node_names=Identity,Identity_1,Identity_2,Identity_3 {f_pb} {f}' f'--output_node_names=Identity,Identity_1,Identity_2,Identity_3 {f_pb} {f}'
subprocess.run(cmd.split()) subprocess.run(cmd.split())
json = Path(f_json).read_text()
with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order
subst = re.sub( subst = re.sub(
r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, ' r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
@ -443,12 +561,11 @@ def export_tfjs(file, prefix=colorstr('TensorFlow.js:')):
r'"Identity.?.?": {"name": "Identity.?.?"}}}', r'{"outputs": {"Identity": {"name": "Identity"}, ' r'"Identity.?.?": {"name": "Identity.?.?"}}}', r'{"outputs": {"Identity": {"name": "Identity"}, '
r'"Identity_1": {"name": "Identity_1"}, ' r'"Identity_1": {"name": "Identity_1"}, '
r'"Identity_2": {"name": "Identity_2"}, ' r'"Identity_2": {"name": "Identity_2"}, '
r'"Identity_3": {"name": "Identity_3"}}}', json) r'"Identity_3": {"name": "Identity_3"}}}', f_json.read_text())
j.write(subst) j.write(subst)
return f, None return f, None
def _add_tflite_metadata(self, file, num_outputs):
def add_tflite_metadata(file, metadata, num_outputs):
# Add metadata to *.tflite models per https://www.tensorflow.org/lite/models/convert/metadata # Add metadata to *.tflite models per https://www.tensorflow.org/lite/models/convert/metadata
with contextlib.suppress(ImportError): with contextlib.suppress(ImportError):
# check_requirements('tflite_support') # check_requirements('tflite_support')
@ -458,7 +575,7 @@ def add_tflite_metadata(file, metadata, num_outputs):
tmp_file = Path('/tmp/meta.txt') tmp_file = Path('/tmp/meta.txt')
with open(tmp_file, 'w') as meta_f: with open(tmp_file, 'w') as meta_f:
meta_f.write(str(metadata)) meta_f.write(str(self.metadata))
model_meta = _metadata_fb.ModelMetadataT() model_meta = _metadata_fb.ModelMetadataT()
label_file = _metadata_fb.AssociatedFileT() label_file = _metadata_fb.AssociatedFileT()
@ -481,128 +598,26 @@ def add_tflite_metadata(file, metadata, num_outputs):
tmp_file.unlink() tmp_file.unlink()
@smart_inference_mode() @hydra.main(version_base=None, config_path=str(DEFAULT_CONFIG.parent), config_name=DEFAULT_CONFIG.name)
def export_model( def export(cfg):
model, # model cfg.model = cfg.model or "yolov8n.yaml"
file=ROOT / 'yolov8n.pt', cfg.format = cfg.format or "torchscript"
data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' exporter = Exporter(cfg)
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 model = None
warnings.filterwarnings('ignore', category=torch.jit.TracerWarning) # suppress TracerWarning if isinstance(cfg.model, (str, Path)):
warnings.filterwarnings('ignore', category=UserWarning) # suppress shape prim::Constant type missing ONNX warning if Path(cfg.model).suffix == '.yaml':
warnings.filterwarnings('ignore', category=DeprecationWarning) # suppress CoreML np.bool deprecation warning 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)
# 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 if __name__ == "__main__":
f = [str(x) for x in f if x] # filter out '' and None """
if any(f): CLI:
cls, det, seg = (isinstance(model, x) for x in (ClassificationModel, DetectionModel, SegmentationModel)) # type yolo mode=export model=yolov8n.yaml format=onnx
det &= not seg # segmentation models inherit from SegmentationModel(DetectionModel) """
dir = Path('segment' if seg else 'classify' if cls else '') export()
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

37
ultralytics/yolo/engine/model.py
Unescape View File

@ -5,7 +5,7 @@ import torch
from ultralytics import yolo # noqa required for python usage from ultralytics import yolo # noqa required for python usage
from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, attempt_load_weights from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, attempt_load_weights
from ultralytics.yolo.configs import get_config 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 import DEFAULT_CONFIG, HELP_MSG, LOGGER
from ultralytics.yolo.utils.checks import check_yaml from ultralytics.yolo.utils.checks import check_yaml
from ultralytics.yolo.utils.files import yaml_load from ultralytics.yolo.utils.files import yaml_load
@ -164,7 +164,7 @@ class YOLO:
validator(model=self.model) validator(model=self.model)
@smart_inference_mode() @smart_inference_mode()
def export(self, format='', save_dir='', **kwargs): def export(self, **kwargs):
""" """
Export model. Export model.
@ -177,36 +177,9 @@ class YOLO:
overrides.update(kwargs) overrides.update(kwargs)
args = get_config(config=DEFAULT_CONFIG, overrides=overrides) args = get_config(config=DEFAULT_CONFIG, overrides=overrides)
args.task = self.task args.task = self.task
args.format = format
exporter = Exporter(overrides=overrides)
file = self.ckpt or Path(Path(self.cfg).name) exporter(model=self.model)
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
)
def train(self, **kwargs): def train(self, **kwargs):
""" """

8
ultralytics/yolo/engine/predictor.py
Unescape View File

@ -16,14 +16,14 @@ Usage - formats:
$ yolo task=... mode=predict --weights yolov8n.pt # PyTorch $ yolo task=... mode=predict --weights yolov8n.pt # PyTorch
yolov8n.torchscript # TorchScript yolov8n.torchscript # TorchScript
yolov8n.onnx # ONNX Runtime or OpenCV DNN with --dnn yolov8n.onnx # ONNX Runtime or OpenCV DNN with --dnn
yolov5s_openvino_model # OpenVINO yolov8n_openvino_model # OpenVINO
yolov8n.engine # TensorRT yolov8n.engine # TensorRT
yolov8n.mlmodel # CoreML (macOS-only) yolov8n.mlmodel # CoreML (macOS-only)
yolov5s_saved_model # TensorFlow SavedModel yolov8n_saved_model # TensorFlow SavedModel
yolov8n.pb # TensorFlow GraphDef yolov8n.pb # TensorFlow GraphDef
yolov8n.tflite # TensorFlow Lite yolov8n.tflite # TensorFlow Lite
yolov5s_edgetpu.tflite # TensorFlow Edge TPU yolov8n_edgetpu.tflite # TensorFlow Edge TPU
yolov5s_paddle_model # PaddlePaddle yolov8n_paddle_model # PaddlePaddle
""" """
import platform import platform
from pathlib import Path from pathlib import Path

20
ultralytics/yolo/utils/__init__.py
Unescape View File

@ -25,14 +25,12 @@ TQDM_BAR_FORMAT = '{l_bar}{bar:10}{r_bar}' # tqdm bar format
LOGGING_NAME = 'yolov5' LOGGING_NAME = 'yolov5'
HELP_MSG = \ HELP_MSG = \
""" """
Please refer to below Usage examples for help running YOLOv8 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
Install: Install:
pip install ultralytics pip install ultralytics
Python usage: Python SDK:
from ultralytics import YOLO from ultralytics import YOLO
model = YOLO.new('yolov8n.yaml') # create a new model from scratch model = YOLO.new('yolov8n.yaml') # create a new model from scratch
@ -42,12 +40,15 @@ HELP_MSG = \
results = model.predict(source='bus.jpg') results = model.predict(source='bus.jpg')
success = model.export(format='onnx') success = model.export(format='onnx')
CLI usage: CLI:
yolo task=detect mode=train model=yolov8n.yaml ... yolo task=detect mode=train model=yolov8n.yaml args...
classify predict yolov8n-cls.yaml classify predict yolov8n-cls.yaml args...
segment val yolov8n-seg.yaml 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 # Settings
@ -56,7 +57,6 @@ HELP_MSG = \
pd.options.display.max_columns = 10 pd.options.display.max_columns = 10
cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader) cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader)
os.environ['NUMEXPR_MAX_THREADS'] = str(NUM_THREADS) # NumExpr max threads 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(): def is_colab():

4
ultralytics/yolo/utils/callbacks/clearml.py
Unescape View File

@ -36,8 +36,8 @@ def on_val_end(trainer):
if trainer.epoch == 0: if trainer.epoch == 0:
model_info = { model_info = {
"Parameters": get_num_params(trainer.model), "Parameters": get_num_params(trainer.model),
"GFLOPs": round(get_flops(trainer.model), 1), "GFLOPs": round(get_flops(trainer.model), 3),
"Inference speed (ms/img)": round(trainer.validator.speed[1], 1)} "Inference speed (ms/img)": round(trainer.validator.speed[1], 3)}
Task.current_task().connect(model_info, name='Model') Task.current_task().connect(model_info, name='Model')

4
ultralytics/yolo/utils/callbacks/wb.py
Unescape View File

@ -19,8 +19,8 @@ def on_val_end(trainer):
if trainer.epoch == 0: if trainer.epoch == 0:
model_info = { model_info = {
"model/parameters": get_num_params(trainer.model), "model/parameters": get_num_params(trainer.model),
"model/GFLOPs": round(get_flops(trainer.model), 1), "model/GFLOPs": round(get_flops(trainer.model), 3),
"model/speed(ms)": round(trainer.validator.speed[1], 1)} "model/speed(ms)": round(trainer.validator.speed[1], 3)}
wandb.run.log(model_info, step=trainer.epoch + 1) wandb.run.log(model_info, step=trainer.epoch + 1)

Write Preview
Loading…
Cancel
Save