diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index a42e0bb431..eb41f841ae 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -1054,11 +1054,10 @@ class Exporter: @try_export def export_mct(self, prefix=colorstr("Sony MCT:")): - check_requirements(["mct-nightly", "sony-custom-layers[torch]"]) + check_requirements(["model_compression_toolkit==2.1.0", "sony-custom-layers[torch]"]) import model_compression_toolkit as mct import onnx - from model_compression_toolkit.core import BitWidthConfig - from model_compression_toolkit.core.common.network_editors import NodeNameScopeFilter + import subprocess from model_compression_toolkit.core.pytorch.pytorch_device_config import get_working_device, set_working_device from sony_custom_layers.pytorch.object_detection.nms import multiclass_nms @@ -1109,25 +1108,12 @@ class Exporter: yield [img] tpc = mct.get_target_platform_capabilities( - fw_name="pytorch", target_platform_name="imx500", target_platform_version="v3" - ) - - # Configure MCT manually for specific layers - bit_cfg = BitWidthConfig() - bit_cfg.set_manual_activation_bit_width( - [ - NodeNameScopeFilter("mul"), - NodeNameScopeFilter("sub"), - NodeNameScopeFilter("add_6"), - NodeNameScopeFilter("cat_17"), - ], - 16, + fw_name="pytorch", target_platform_name="imx500", target_platform_version="v1" ) config = mct.core.CoreConfig( mixed_precision_config=mct.core.MixedPrecisionQuantizationConfig(num_of_images=10), - quantization_config=mct.core.QuantizationConfig(concat_threshold_update=True), - bit_width_config=bit_cfg, + quantization_config=mct.core.QuantizationConfig(concat_threshold_update=True) ) resource_utilization = mct.core.ResourceUtilization(weights_memory=3146176 * 0.76) @@ -1141,7 +1127,6 @@ class Exporter: core_config=config, target_platform_capabilities=tpc, ) - print("Quantized model is ready") else: gptq_config = mct.gptq.get_pytorch_gptq_config(n_epochs=1000, use_hessian_based_weights=False) @@ -1157,8 +1142,6 @@ class Exporter: target_platform_capabilities=tpc, ) - print("Quantized-PTQ model is ready") - if self.args.nms: # Define PostProcess params score_threshold = 0.001 @@ -1171,7 +1154,7 @@ class Exporter: iou_threshold=iou_threshold, max_detections=max_detections, ).to(device=get_working_device()) - + f = Path(str(self.file).replace(self.file.suffix, "_mct_model.onnx")) # js dir mct.exporter.pytorch_export_model(model=quant_model, save_model_path=f, repr_dataset=representative_dataset_gen) @@ -1181,6 +1164,19 @@ class Exporter: meta.key, meta.value = k, str(v) onnx.save(model_onnx, f) + + if not LINUX: + LOGGER.warning(f"{prefix} WARNING ⚠️ MCT imx500-converter is only supported on Linux.") + else: + check_requirements("imx500-converter[pt]==3.14.1") + try: + subprocess.run(["java", "--version"], check=True) + except FileNotFoundError: + LOGGER.error("Java 17 is required for the imx500 conversion. \n Please install Java with: \n sudo apt install openjdk-17-jdk openjdk-17-jre") + return None + + subprocess.run(["imxconv-pt", "-i", "yolov8n_mct_model.onnx", "-o", "yolov8n_imx500_model"], check=True) + return f, None def _add_tflite_metadata(self, file): diff --git a/ultralytics/nn/autobackend.py b/ultralytics/nn/autobackend.py index 8e36d54450..755df474a8 100644 --- a/ultralytics/nn/autobackend.py +++ b/ultralytics/nn/autobackend.py @@ -178,7 +178,7 @@ class AutoBackend(nn.Module): providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if cuda else ["CPUExecutionProvider"] if mct: - check_requirements(["mct-nightly", "sony-custom-layers[torch]"]) + check_requirements(["model_compression_toolkit==2.1.0", "sony-custom-layers[torch]", "onnxruntime-extensions"]) LOGGER.info(f"Loading {w} for ONNX MCT quantization inference...") import mct_quantizers as mctq from sony_custom_layers.pytorch.object_detection import nms_ort # noqa