Improved OpenVINO quantization code (#12614)
* Improved OpenVINO quantization code * Auto-format by Ultralytics actions --------- Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
1 parent
9abbef522f
commit
de64179cd4
1 file changed
+11
-20
@@ -226,28 +226,19 @@ def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:
|
|||||||
from openvino.tools import mo # noqa
|
from openvino.tools import mo # noqa
|
||||||
|
|
||||||
LOGGER.info(f"\n{prefix} starting export with openvino {ov.__version__}...")
|
LOGGER.info(f"\n{prefix} starting export with openvino {ov.__version__}...")
|
||||||
f = str(file).replace(file.suffix, f"_openvino_model{os.sep}")
|
f = str(file).replace(file.suffix, f"_{'int8_' if int8 else ''}openvino_model{os.sep}")
|
||||||
f_onnx = file.with_suffix(".onnx")
|
f_onnx = file.with_suffix(".onnx")
|
||||||
f_ov = str(Path(f) / file.with_suffix(".xml").name)
|
f_ov = str(Path(f) / file.with_suffix(".xml").name)
|
||||||
|
|
||||||
|
ov_model = mo.convert_model(f_onnx, model_name=file.stem, framework="onnx", compress_to_fp16=half) # export
|
||||||
|
|
||||||
if int8:
|
if int8:
|
||||||
check_requirements("nncf>=2.4.0") # requires at least version 2.4.0 to use the post-training quantization
|
check_requirements("nncf>=2.5.0") # requires at least version 2.5.0 to use the post-training quantization
|
||||||
import nncf
|
import nncf
|
||||||
import numpy as np
|
import numpy as np
|
||||||
from openvino.runtime import Core
|
|
||||||
|
|
||||||
from utils.dataloaders import create_dataloader
|
from utils.dataloaders import create_dataloader
|
||||||
|
|
||||||
core = Core()
|
|
||||||
onnx_model = core.read_model(f_onnx) # export
|
|
||||||
|
|
||||||
def prepare_input_tensor(image: np.ndarray):
|
|
||||||
input_tensor = image.astype(np.float32) # uint8 to fp16/32
|
|
||||||
input_tensor /= 255.0 # 0 - 255 to 0.0 - 1.0
|
|
||||||
|
|
||||||
if input_tensor.ndim == 3:
|
|
||||||
input_tensor = np.expand_dims(input_tensor, 0)
|
|
||||||
return input_tensor
|
|
||||||
|
|
||||||
def gen_dataloader(yaml_path, task="train", imgsz=640, workers=4):
|
def gen_dataloader(yaml_path, task="train", imgsz=640, workers=4):
|
||||||
data_yaml = check_yaml(yaml_path)
|
data_yaml = check_yaml(yaml_path)
|
||||||
data = check_dataset(data_yaml)
|
data = check_dataset(data_yaml)
|
||||||
@@ -268,15 +259,15 @@ def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:
|
|||||||
Returns:
|
Returns:
|
||||||
input_tensor: Input data for quantization
|
input_tensor: Input data for quantization
|
||||||
"""
|
"""
|
||||||
img = data_item[0].numpy()
|
assert data_item[0].dtype == torch.uint8, "input image must be uint8 for the quantization preprocessing"
|
||||||
input_tensor = prepare_input_tensor(img)
|
|
||||||
return input_tensor
|
img = data_item[0].numpy().astype(np.float32) # uint8 to fp16/32
|
||||||
|
img /= 255.0 # 0 - 255 to 0.0 - 1.0
|
||||||
|
return np.expand_dims(img, 0) if img.ndim == 3 else img
|
||||||
|
|
||||||
ds = gen_dataloader(data)
|
ds = gen_dataloader(data)
|
||||||
quantization_dataset = nncf.Dataset(ds, transform_fn)
|
quantization_dataset = nncf.Dataset(ds, transform_fn)
|
||||||
ov_model = nncf.quantize(onnx_model, quantization_dataset, preset=nncf.QuantizationPreset.MIXED)
|
ov_model = nncf.quantize(ov_model, quantization_dataset, preset=nncf.QuantizationPreset.MIXED)
|
||||||
else:
|
|
||||||
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
|
ov.serialize(ov_model, f_ov) # save
|
||||||
yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata) # add metadata.yaml
|
yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata) # add metadata.yaml
|
||||||
|
|||||||
Reference in new issue
Block a user