Add `--half` support for OpenVINO exports (#7615)
Browse files* feature(export): add half support for openvino models
* Update export.py
* Update export.py
Co-authored-by: Glenn Jocher <[email protected]>
export.py
CHANGED
@@ -168,7 +168,7 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
|
|
168 |
LOGGER.info(f'{prefix} export failure: {e}')
|
169 |
|
170 |
|
171 |
-
def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')):
|
172 |
# YOLOv5 OpenVINO export
|
173 |
try:
|
174 |
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:')):
|
|
177 |
LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...')
|
178 |
f = str(file).replace('.pt', '_openvino_model' + os.sep)
|
179 |
|
180 |
-
cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f}"
|
181 |
subprocess.check_output(cmd, shell=True)
|
182 |
|
183 |
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
@@ -477,7 +477,7 @@ def run(
|
|
477 |
# Load PyTorch model
|
478 |
device = select_device(device)
|
479 |
if half:
|
480 |
-
assert device.type != 'cpu' or coreml, '--half only compatible with GPU export, i.e. use --device 0'
|
481 |
model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model
|
482 |
nc, names = model.nc, model.names # number of classes, class names
|
483 |
|
@@ -491,7 +491,7 @@ def run(
|
|
491 |
im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection
|
492 |
|
493 |
# Update model
|
494 |
-
if half and not coreml:
|
495 |
im, model = im.half(), model.half() # to FP16
|
496 |
model.train() if train else model.eval() # training mode = no Detect() layer grid construction
|
497 |
for k, m in model.named_modules():
|
@@ -515,7 +515,7 @@ def run(
|
|
515 |
if onnx or xml: # OpenVINO requires ONNX
|
516 |
f[2] = export_onnx(model, im, file, opset, train, dynamic, simplify)
|
517 |
if xml: # OpenVINO
|
518 |
-
f[3] = export_openvino(model, im, file)
|
519 |
if coreml:
|
520 |
_, f[4] = export_coreml(model, im, file, int8, half)
|
521 |
|
|
|
168 |
LOGGER.info(f'{prefix} export failure: {e}')
|
169 |
|
170 |
|
171 |
+
def export_openvino(model, im, file, half, prefix=colorstr('OpenVINO:')):
|
172 |
# YOLOv5 OpenVINO export
|
173 |
try:
|
174 |
check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/
|
|
|
177 |
LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...')
|
178 |
f = str(file).replace('.pt', '_openvino_model' + os.sep)
|
179 |
|
180 |
+
cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} --data_type {'FP16' if half else 'FP32'}"
|
181 |
subprocess.check_output(cmd, shell=True)
|
182 |
|
183 |
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
|
|
477 |
# Load PyTorch model
|
478 |
device = select_device(device)
|
479 |
if half:
|
480 |
+
assert device.type != 'cpu' or coreml or xml, '--half only compatible with GPU export, i.e. use --device 0'
|
481 |
model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model
|
482 |
nc, names = model.nc, model.names # number of classes, class names
|
483 |
|
|
|
491 |
im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection
|
492 |
|
493 |
# Update model
|
494 |
+
if half and not (coreml or xml):
|
495 |
im, model = im.half(), model.half() # to FP16
|
496 |
model.train() if train else model.eval() # training mode = no Detect() layer grid construction
|
497 |
for k, m in model.named_modules():
|
|
|
515 |
if onnx or xml: # OpenVINO requires ONNX
|
516 |
f[2] = export_onnx(model, im, file, opset, train, dynamic, simplify)
|
517 |
if xml: # OpenVINO
|
518 |
+
f[3] = export_openvino(model, im, file, half)
|
519 |
if coreml:
|
520 |
_, f[4] = export_coreml(model, im, file, int8, half)
|
521 |
|