From 1771bcdc1a7540ff158ec3df2798d4e66402e561 Mon Sep 17 00:00:00 2001 From: Duncan Moss Date: Wed, 27 Apr 2022 16:37:39 -0700 Subject: [PATCH 1/3] feature(export): add half support for openvino models --- export.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/export.py b/export.py index 3b85dfc27f27..6fd22f2729bd 100644 --- a/export.py +++ b/export.py @@ -168,7 +168,7 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst LOGGER.info(f'{prefix} export failure: {e}') -def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')): +def export_openvino(model, im, file, half, prefix=colorstr('OpenVINO:')): # YOLOv5 OpenVINO export try: check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/ @@ -177,7 +177,7 @@ def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')): LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...') f = str(file).replace('.pt', '_openvino_model' + os.sep) - cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f}" + cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} {'--data_type FP16' if half else ''}" subprocess.check_output(cmd, shell=True) LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') @@ -477,7 +477,7 @@ def run( # 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 device.type != 'cpu' or coreml or xml, '--half only compatible with GPU and OpenVINO export, i.e. use --device 0' model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model nc, names = model.nc, model.names # number of classes, class names @@ -491,7 +491,7 @@ def run( im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection # Update model - if half and not coreml: + if half and not (coreml or xml): im, model = im.half(), model.half() # to FP16 model.train() if train else model.eval() # training mode = no Detect() layer grid construction for k, m in model.named_modules(): @@ -515,7 +515,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(model, im, file) + f[3] = export_openvino(model, im, file, half) if coreml: _, f[4] = export_coreml(model, im, file, int8, half) From 9f65617e92ef42231fcef1dd3db865c0d001fc1f Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 27 Apr 2022 17:23:53 -0700 Subject: [PATCH 2/3] Update export.py --- export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/export.py b/export.py index 6fd22f2729bd..0b98f572db4a 100644 --- a/export.py +++ b/export.py @@ -177,7 +177,7 @@ def export_openvino(model, im, file, half, prefix=colorstr('OpenVINO:')): LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...') f = str(file).replace('.pt', '_openvino_model' + os.sep) - cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} {'--data_type FP16' if half else ''}" + cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} --data_type {'FP16' if half else 'FP32'}" subprocess.check_output(cmd, shell=True) LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') From e1e6737d45349959f83572d1c296bca88928aed2 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 27 Apr 2022 17:27:49 -0700 Subject: [PATCH 3/3] Update export.py --- export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/export.py b/export.py index 0b98f572db4a..d4191d80b410 100644 --- a/export.py +++ b/export.py @@ -477,7 +477,7 @@ def run( # Load PyTorch model device = select_device(device) if half: - assert device.type != 'cpu' or coreml or xml, '--half only compatible with GPU and OpenVINO export, i.e. use --device 0' + assert device.type != 'cpu' or coreml or xml, '--half only compatible with GPU export, i.e. use --device 0' model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model nc, names = model.nc, model.names # number of classes, class names