Skip to content

Commit

Permalink
Fix OpenVINO export (#11666)
Browse files Browse the repository at this point in the history
* Fix OpenVINO export

Resolves #11645

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update export.py

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

---------

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
glenn-jocher and pre-commit-ci[bot] committed Jun 6, 2023
1 parent 76ea9ed commit 89c3040
Showing 1 changed file with 10 additions and 13 deletions.
23 changes: 10 additions & 13 deletions export.py
Original file line number Diff line number Diff line change
Expand Up @@ -207,21 +207,18 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('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
check_requirements('openvino-dev>=2022.3') # requires openvino-dev: https://pypi.org/project/openvino-dev/
import openvino.runtime as ov # noqa
from openvino.tools import mo # noqa

LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...')
f = str(file).replace('.pt', f'_openvino_model{os.sep}')
LOGGER.info(f'\n{prefix} starting export with openvino {ov.__version__}...')
f = str(file).replace(file.suffix, f'_openvino_model{os.sep}')
f_onnx = file.with_suffix('.onnx')
f_ov = str(Path(f) / file.with_suffix('.xml').name)

args = [
'mo',
'--input_model',
str(file.with_suffix('.onnx')),
'--output_dir',
f,
'--data_type',
('FP16' if half else 'FP32'),]
subprocess.run(args, check=True, env=os.environ) # export
ov_model = mo.convert_model(f_onnx, model_name=file.stem, framework='onnx', compress_to_fp16=half) # export

ov.serialize(ov_model, f_ov) # save
yaml_save(Path(f) / file.with_suffix('.yaml').name, metadata) # add metadata.yaml
return f, None

Expand Down

0 comments on commit 89c3040

Please sign in to comment.