Add OpenVINO metadata to export (#7947)
* Write .yaml file when exporting model to openvino Write .yaml file automatically when exporting model to openvino to be used during inference * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update export.py * Update export.py * Load metadata on inference * Update common.py Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
2 files changed
+24
-12
No files matched your search
@@ -54,6 +54,7 @@ from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
import yaml
|
||||
from torch.utils.mobile_optimizer import optimize_for_mobile
|
||||
|
||||
FILE = Path(__file__).resolve()
|
||||
@@ -168,7 +169,7 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
|
||||
LOGGER.info(f'{prefix} export failure: {e}')
|
||||
|
||||
|
||||
def export_openvino(file, half, prefix=colorstr('OpenVINO:')):
|
||||
def export_openvino(model, file, half, prefix=colorstr('OpenVINO:')):
|
||||
# YOLOv5 OpenVINO export
|
||||
try:
|
||||
check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/
|
||||
@@ -178,7 +179,9 @@ def export_openvino(file, half, prefix=colorstr('OpenVINO:')):
|
||||
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.check_output(cmd.split())
|
||||
subprocess.check_output(cmd.split()) # export
|
||||
with open(Path(f) / 'meta.yaml', 'w') as g:
|
||||
yaml.dump({'stride': int(max(model.stride)), 'names': model.names}, g) # add metadata.yaml
|
||||
|
||||
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
||||
return f
|
||||
@@ -520,7 +523,7 @@ def run(
|
||||
if onnx or xml: # OpenVINO requires ONNX
|
||||
f[2] = export_onnx(model, im, file, opset, train, dynamic, simplify)
|
||||
if xml: # OpenVINO
|
||||
f[3] = export_openvino(file, half)
|
||||
f[3] = export_openvino(model, file, half)
|
||||
if coreml:
|
||||
_, f[4] = export_coreml(model, im, file, int8, half)
|
||||
|
||||
|
||||
Reference in new issue
Block a user