From d9de9e9721e03cc4cacc5e46dbb23e1367462214 Mon Sep 17 00:00:00 2001 From: Duncan Moss Date: Wed, 27 Apr 2022 17:29:59 -0700 Subject: [PATCH] Add `--half` support for OpenVINO exports (#7615) * feature(export): add half support for openvino models * Update export.py * Update export.py Co-authored-by: Glenn Jocher --- export.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/export.py b/export.py index 4ced8b155c52..f7509b5c509c 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 'FP32'}" subprocess.check_output(cmd, shell=True) LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') @@ -596,7 +596,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 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 @@ -610,7 +610,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(): @@ -634,7 +634,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: nb = shape[1] _, f[4] = export_coreml(model, im, file, nb, nc, names, conf_thres, iou_thres, int8, half)