mct-nightly 2.1.0.20240807.445__py3-none-any.whl → 2.1.0.20240808.431__py3-none-any.whl

model_compression_toolkit/metadata.py

@@ -13,8 +13,31 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Dict
-from model_compression_toolkit.constants import MCT_VERSION, TPC_VERSION
+from typing import Dict, Any
+from model_compression_toolkit.constants import MCT_VERSION, TPC_VERSION, OPERATORS_SCHEDULING, FUSED_NODES_MAPPING, \
+    CUTS, MAX_CUT, OP_ORDER, OP_RECORD, SHAPE, NODE_OUTPUT_INDEX, NODE_NAME, TOTAL_SIZE, MEM_ELEMENTS
+from model_compression_toolkit.core.common.graph.memory_graph.compute_graph_max_cut import SchedulerInfo
+from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
+
+
+def create_model_metadata(tpc: TargetPlatformCapabilities,
+                          scheduling_info: SchedulerInfo = None) -> Dict:
+    """
+    Creates and returns a metadata dictionary for the model, including version information
+    and optional scheduling information.
+
+    Args:
+        tpc: A TPC object to get the version.
+        scheduling_info: An object containing scheduling details and metadata. Default is None.
+
+    Returns:
+        Dict: A dictionary containing the model's version information and optional scheduling information.
+    """
+    _metadata = get_versions_dict(tpc)
+    if scheduling_info:
+        scheduler_metadata = get_scheduler_metadata(scheduler_info=scheduling_info)
+        _metadata['scheduling_info'] = scheduler_metadata
+    return _metadata
 
 
 def get_versions_dict(tpc) -> Dict:
@@ -27,3 +50,39 @@ def get_versions_dict(tpc) -> Dict:
     from model_compression_toolkit import __version__ as mct_version
     tpc_version = f'{tpc.name}.{tpc.version}'
     return {MCT_VERSION: mct_version, TPC_VERSION: tpc_version}
+
+
+def get_scheduler_metadata(scheduler_info: SchedulerInfo) -> Dict[str, Any]:
+    """
+    Extracts and returns metadata from SchedulerInfo.
+
+    Args:
+        scheduler_info (SchedulerInfo): The scheduler information object containing scheduling details like cuts and
+        fusing mapping.
+
+    Returns:
+        Dict[str, Any]: A dictionary containing extracted metadata, including schedule, maximum cut,
+        cuts information, and fused nodes mapping.
+    """
+    scheduler_metadata = {
+        OPERATORS_SCHEDULING: [str(layer) for layer in scheduler_info.operators_scheduling],
+        MAX_CUT: scheduler_info.max_cut,
+        CUTS: [
+            {
+                OP_ORDER: [op.name for op in cut.op_order],
+                OP_RECORD: [op.name for op in cut.op_record],
+                MEM_ELEMENTS: [
+                    {
+                        SHAPE: list(tensor.shape),
+                        NODE_NAME: tensor.node_name,
+                        TOTAL_SIZE: float(tensor.total_size),
+                        NODE_OUTPUT_INDEX: tensor.node_output_index
+                    }
+                    for tensor in cut.mem_elements.elements
+                ]
+            }
+            for cut in scheduler_info.cuts
+        ],
+        FUSED_NODES_MAPPING: scheduler_info.fused_nodes_mapping
+    }
+    return scheduler_metadata

model_compression_toolkit/ptq/keras/quantization_facade.py

@@ -28,7 +28,7 @@ from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quant
 from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import TargetPlatformCapabilities
 from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.ptq.runner import ptq_runner
-from model_compression_toolkit.metadata import get_versions_dict
+from model_compression_toolkit.metadata import create_model_metadata
 
 if FOUND_TF:
     from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
@@ -134,14 +134,14 @@ if FOUND_TF:
         fw_impl = KerasImplementation()
 
         # Ignore returned hessian service as PTQ does not use it
-        tg, bit_widths_config, _ = core_runner(in_model=in_model,
-                                               representative_data_gen=representative_data_gen,
-                                               core_config=core_config,
-                                               fw_info=fw_info,
-                                               fw_impl=fw_impl,
-                                               tpc=target_platform_capabilities,
-                                               target_resource_utilization=target_resource_utilization,
-                                               tb_w=tb_w)
+        tg, bit_widths_config, _, scheduling_info = core_runner(in_model=in_model,
+                                                                representative_data_gen=representative_data_gen,
+                                                                core_config=core_config,
+                                                                fw_info=fw_info,
+                                                                fw_impl=fw_impl,
+                                                                tpc=target_platform_capabilities,
+                                                                target_resource_utilization=target_resource_utilization,
+                                                                tb_w=tb_w)
 
         # At this point, tg is a graph that went through substitutions (such as BN folding) and is
         # ready for quantization (namely, it holds quantization params, etc.) but the weights are
@@ -168,7 +168,9 @@ if FOUND_TF:
 
         exportable_model, user_info = get_exportable_keras_model(graph_with_stats_correction)
         if target_platform_capabilities.tp_model.add_metadata:
-            exportable_model = add_metadata(exportable_model, get_versions_dict(target_platform_capabilities))
+            exportable_model = add_metadata(exportable_model,
+                                            create_model_metadata(tpc=target_platform_capabilities,
+                                                                  scheduling_info=scheduling_info))
         return exportable_model, user_info
 
 

model_compression_toolkit/ptq/pytorch/quantization_facade.py

@@ -16,7 +16,6 @@ import copy
 
 from typing import Callable
 
-from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import PYTORCH, FOUND_TORCH
@@ -29,8 +28,7 @@ from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.ptq.runner import ptq_runner
 from model_compression_toolkit.core.analyzer import analyzer_model_quantization
 from model_compression_toolkit.core.common.quantization.quantize_graph_weights import quantize_graph_weights
-from model_compression_toolkit.metadata import get_versions_dict
-
+from model_compression_toolkit.metadata import create_model_metadata
 
 if FOUND_TORCH:
     from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
@@ -109,14 +107,14 @@ if FOUND_TORCH:
         fw_impl = PytorchImplementation()
 
         # Ignore hessian info service as it is not used here yet.
-        tg, bit_widths_config, _ = core_runner(in_model=in_module,
-                                               representative_data_gen=representative_data_gen,
-                                               core_config=core_config,
-                                               fw_info=fw_info,
-                                               fw_impl=fw_impl,
-                                               tpc=target_platform_capabilities,
-                                               target_resource_utilization=target_resource_utilization,
-                                               tb_w=tb_w)
+        tg, bit_widths_config, _, scheduling_info = core_runner(in_model=in_module,
+                                                                representative_data_gen=representative_data_gen,
+                                                                core_config=core_config,
+                                                                fw_info=fw_info,
+                                                                fw_impl=fw_impl,
+                                                                tpc=target_platform_capabilities,
+                                                                target_resource_utilization=target_resource_utilization,
+                                                                tb_w=tb_w)
 
         # At this point, tg is a graph that went through substitutions (such as BN folding) and is
         # ready for quantization (namely, it holds quantization params, etc.) but the weights are
@@ -143,7 +141,9 @@ if FOUND_TORCH:
 
         exportable_model, user_info = get_exportable_pytorch_model(graph_with_stats_correction)
         if target_platform_capabilities.tp_model.add_metadata:
-            exportable_model = add_metadata(exportable_model, get_versions_dict(target_platform_capabilities))
+            exportable_model = add_metadata(exportable_model,
+                                            create_model_metadata(tpc=target_platform_capabilities,
+                                                                  scheduling_info=scheduling_info))
         return exportable_model, user_info
 
 

model_compression_toolkit/qat/keras/quantization_facade.py

@@ -187,14 +187,14 @@ if FOUND_TF:
         fw_impl = KerasImplementation()
 
         # Ignore hessian service since is not used in QAT at the moment
-        tg, bit_widths_config, _ = core_runner(in_model=in_model,
-                                               representative_data_gen=representative_data_gen,
-                                               core_config=core_config,
-                                               fw_info=DEFAULT_KERAS_INFO,
-                                               fw_impl=fw_impl,
-                                               tpc=target_platform_capabilities,
-                                               target_resource_utilization=target_resource_utilization,
-                                               tb_w=tb_w)
+        tg, bit_widths_config, _, _ = core_runner(in_model=in_model,
+                                                  representative_data_gen=representative_data_gen,
+                                                  core_config=core_config,
+                                                  fw_info=DEFAULT_KERAS_INFO,
+                                                  fw_impl=fw_impl,
+                                                  tpc=target_platform_capabilities,
+                                                  target_resource_utilization=target_resource_utilization,
+                                                  tb_w=tb_w)
 
         tg = ptq_runner(tg, representative_data_gen, core_config, DEFAULT_KERAS_INFO, fw_impl, tb_w)
 

model_compression_toolkit/qat/pytorch/quantization_facade.py

@@ -154,14 +154,14 @@ if FOUND_TORCH:
         fw_impl = PytorchImplementation()
 
         # Ignore hessian scores service as we do not use it here
-        tg, bit_widths_config, _ = core_runner(in_model=in_model,
-                                               representative_data_gen=representative_data_gen,
-                                               core_config=core_config,
-                                               fw_info=DEFAULT_PYTORCH_INFO,
-                                               fw_impl=fw_impl,
-                                               tpc=target_platform_capabilities,
-                                               target_resource_utilization=target_resource_utilization,
-                                               tb_w=tb_w)
+        tg, bit_widths_config, _, _ = core_runner(in_model=in_model,
+                                                  representative_data_gen=representative_data_gen,
+                                                  core_config=core_config,
+                                                  fw_info=DEFAULT_PYTORCH_INFO,
+                                                  fw_impl=fw_impl,
+                                                  tpc=target_platform_capabilities,
+                                                  target_resource_utilization=target_resource_utilization,
+                                                  tb_w=tb_w)
 
         tg = ptq_runner(tg, representative_data_gen, core_config, DEFAULT_PYTORCH_INFO, fw_impl, tb_w)
 

model_compression_toolkit/target_platform_capabilities/tpc_models/get_target_platform_capabilities.py

@@ -16,17 +16,17 @@
 from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
 
 from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.target_platform_capabilities import \
-    tpc_dict as imx500_tpc_dict
+    get_tpc_dict_by_fw as get_imx500_tpc
 from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.target_platform_capabilities import \
-    tpc_dict as tflite_tpc_dict
+    get_tpc_dict_by_fw as get_tflite_tpc
 from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.target_platform_capabilities import \
-    tpc_dict as qnnpack_tpc_dict
+    get_tpc_dict_by_fw as get_qnnpack_tpc
 from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL, IMX500_TP_MODEL, TFLITE_TP_MODEL, QNNPACK_TP_MODEL,  LATEST
 
-tpc_dict = {DEFAULT_TP_MODEL: imx500_tpc_dict,
-            IMX500_TP_MODEL: imx500_tpc_dict,
-            TFLITE_TP_MODEL: tflite_tpc_dict,
-            QNNPACK_TP_MODEL: qnnpack_tpc_dict}
+tpc_dict = {DEFAULT_TP_MODEL: get_imx500_tpc,
+            IMX500_TP_MODEL: get_imx500_tpc,
+            TFLITE_TP_MODEL: get_tflite_tpc,
+            QNNPACK_TP_MODEL: get_qnnpack_tpc}
 
 
 def get_target_platform_capabilities(fw_name: str,
@@ -47,13 +47,10 @@ def get_target_platform_capabilities(fw_name: str,
     """
     assert target_platform_name in tpc_dict, f'Target platform {target_platform_name} is not defined!'
     fw_tpc = tpc_dict.get(target_platform_name)
-    assert fw_name in fw_tpc, f'Framework {fw_name} is not supported in {target_platform_name}. Please make sure the relevant ' \
-                              f'packages are installed when using MCT for optimizing a {fw_name} model. ' \
-                              f'For Tensorflow, please install tensorflow. ' \
-                              f'For PyTorch, please install torch.'
-    tpc_versions = fw_tpc.get(fw_name)
+    tpc_versions = fw_tpc(fw_name)
     if target_platform_version is None:
         target_platform_version = LATEST
     else:
-        assert target_platform_version in tpc_versions, f'TPC version {target_platform_version} is not supported for framework {fw_name}.'
+        assert target_platform_version in tpc_versions, (f'TPC version {target_platform_version} is not supported for '
+                                                         f'framework {fw_name}.')
     return tpc_versions[target_platform_version]()

model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/target_platform_capabilities.py

@@ -12,61 +12,74 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+from model_compression_toolkit.logger import Logger
 
 from model_compression_toolkit.constants import FOUND_TF, FOUND_TORCH, TENSORFLOW, PYTORCH
 from model_compression_toolkit.target_platform_capabilities.constants import LATEST
 
-###############################
-# Build Tensorflow TPC models
-###############################
-keras_tpc_models_dict = None
-if FOUND_TF:
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.latest import get_keras_tpc_latest
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tpc_keras import get_keras_tpc as get_keras_tpc_v1
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1_lut.tpc_keras import get_keras_tpc as get_keras_tpc_v1_lut
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1_pot.tpc_keras import get_keras_tpc as get_keras_tpc_v1_pot
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v2.tpc_keras import get_keras_tpc as get_keras_tpc_v2
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v2_lut.tpc_keras import get_keras_tpc as get_keras_tpc_v2_lut
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v3.tpc_keras import get_keras_tpc as get_keras_tpc_v3
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v3_lut.tpc_keras import get_keras_tpc as get_keras_tpc_v3_lut
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v4.tpc_keras import get_keras_tpc as get_keras_tpc_v4
 
-    # Keras: TPC versioning
-    keras_tpc_models_dict = {'v1': get_keras_tpc_v1,
-                             'v1_lut': get_keras_tpc_v1_lut,
-                             'v1_pot': get_keras_tpc_v1_pot,
-                             'v2': get_keras_tpc_v2,
-                             'v2_lut': get_keras_tpc_v2_lut,
-                             'v3': get_keras_tpc_v3,
-                             'v3_lut': get_keras_tpc_v3_lut,
-                             'v4': get_keras_tpc_v4,
-                             LATEST: get_keras_tpc_latest}
+def get_tpc_dict_by_fw(fw_name):
+    tpc_models_dict = None
+    if fw_name == TENSORFLOW:
+        ###############################
+        # Build Tensorflow TPC models
+        ###############################
+        if FOUND_TF:
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.latest import \
+                get_keras_tpc_latest
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tpc_keras import \
+                get_keras_tpc as get_keras_tpc_v1
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1_lut.tpc_keras import \
+                get_keras_tpc as get_keras_tpc_v1_lut
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1_pot.tpc_keras import \
+                get_keras_tpc as get_keras_tpc_v1_pot
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v2.tpc_keras import \
+                get_keras_tpc as get_keras_tpc_v2
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v2_lut.tpc_keras import \
+                get_keras_tpc as get_keras_tpc_v2_lut
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v3.tpc_keras import \
+                get_keras_tpc as get_keras_tpc_v3
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v3_lut.tpc_keras import \
+                get_keras_tpc as get_keras_tpc_v3_lut
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v4.tpc_keras import \
+                get_keras_tpc as get_keras_tpc_v4
 
-###############################
-# Build Pytorch TPC models
-###############################
-pytorch_tpc_models_dict = None
-if FOUND_TORCH:
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.latest import get_pytorch_tpc_latest
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tpc_pytorch import \
-        get_pytorch_tpc as get_pytorch_tpc_v1
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1_pot.tpc_pytorch import \
-        get_pytorch_tpc as get_pytorch_tpc_v1_pot
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1_lut.tpc_pytorch import \
-        get_pytorch_tpc as get_pytorch_tpc_v1_lut
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v2.tpc_pytorch import \
-        get_pytorch_tpc as get_pytorch_tpc_v2
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v2_lut.tpc_pytorch import \
-        get_pytorch_tpc as get_pytorch_tpc_v2_lut
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v3.tpc_pytorch import \
-        get_pytorch_tpc as get_pytorch_tpc_v3
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v3_lut.tpc_pytorch import \
-        get_pytorch_tpc as get_pytorch_tpc_v3_lut
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v4.tpc_pytorch import \
-        get_pytorch_tpc as get_pytorch_tpc_v4
+            # Keras: TPC versioning
+            tpc_models_dict = {'v1': get_keras_tpc_v1,
+                               'v1_lut': get_keras_tpc_v1_lut,
+                               'v1_pot': get_keras_tpc_v1_pot,
+                               'v2': get_keras_tpc_v2,
+                               'v2_lut': get_keras_tpc_v2_lut,
+                               'v3': get_keras_tpc_v3,
+                               'v3_lut': get_keras_tpc_v3_lut,
+                               'v4': get_keras_tpc_v4,
+                               LATEST: get_keras_tpc_latest}
+    elif fw_name == PYTORCH:
+        ###############################
+        # Build Pytorch TPC models
+        ###############################
+        if FOUND_TORCH:
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.latest import \
+                get_pytorch_tpc_latest
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tpc_pytorch import \
+                get_pytorch_tpc as get_pytorch_tpc_v1
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1_pot.tpc_pytorch import \
+                get_pytorch_tpc as get_pytorch_tpc_v1_pot
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1_lut.tpc_pytorch import \
+                get_pytorch_tpc as get_pytorch_tpc_v1_lut
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v2.tpc_pytorch import \
+                get_pytorch_tpc as get_pytorch_tpc_v2
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v2_lut.tpc_pytorch import \
+                get_pytorch_tpc as get_pytorch_tpc_v2_lut
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v3.tpc_pytorch import \
+                get_pytorch_tpc as get_pytorch_tpc_v3
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v3_lut.tpc_pytorch import \
+                get_pytorch_tpc as get_pytorch_tpc_v3_lut
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v4.tpc_pytorch import \
+                get_pytorch_tpc as get_pytorch_tpc_v4
 
-    # Pytorch: TPC versioning
-    pytorch_tpc_models_dict = {'v1': get_pytorch_tpc_v1,
+            # Pytorch: TPC versioning
+            tpc_models_dict = {'v1': get_pytorch_tpc_v1,
                                'v1_lut': get_pytorch_tpc_v1_lut,
                                'v1_pot': get_pytorch_tpc_v1_pot,
                                'v2': get_pytorch_tpc_v2,
@@ -75,7 +88,10 @@ if FOUND_TORCH:
                                'v3_lut': get_pytorch_tpc_v3_lut,
                                'v4': get_pytorch_tpc_v4,
                                LATEST: get_pytorch_tpc_latest}
-
-tpc_dict = {TENSORFLOW: keras_tpc_models_dict,
-            PYTORCH: pytorch_tpc_models_dict}
-
+    if tpc_models_dict is not None:
+        return tpc_models_dict
+    else:
+        Logger.critical(f'Framework {fw_name} is not supported in imx500 or the relevant packages are not '
+                        f'installed. Please make sure the relevant packages are installed when using MCT for optimizing'
+                        f' a {fw_name} model. For Tensorflow, please install tensorflow. For PyTorch, please install '
+                        f'torch.')  # pragma: no cover

model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/target_platform_capabilities.py

@@ -14,35 +14,41 @@
 # ==============================================================================
 
 from model_compression_toolkit.constants import FOUND_TF, FOUND_TORCH, TENSORFLOW, PYTORCH
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.target_platform_capabilities.constants import LATEST
 
-
-###############################
-# Build Tensorflow TPC models
-###############################
-keras_tpc_models_dict = None
-if FOUND_TF:
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.v1.tpc_keras import get_keras_tpc as get_keras_tpc_v1
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.latest import get_keras_tpc_latest
-
-    # Keras: TPC versioning
-    keras_tpc_models_dict = {'v1': get_keras_tpc_v1,
-                             LATEST: get_keras_tpc_latest}
-
-###############################
-# Build Pytorch TPC models
-###############################
-pytorch_tpc_models_dict = None
-if FOUND_TORCH:
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.v1.tpc_pytorch import \
-        get_pytorch_tpc as get_pytorch_tpc_v1
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.latest import get_pytorch_tpc_latest
-
-    # Pytorch: TPC versioning
-    pytorch_tpc_models_dict = {'v1': get_pytorch_tpc_v1,
+def get_tpc_dict_by_fw(fw_name):
+    tpc_models_dict = None
+    if fw_name == TENSORFLOW:
+        ###############################
+        # Build Tensorflow TPC models
+        ###############################
+        if FOUND_TF:
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.v1.tpc_keras import \
+                get_keras_tpc as get_keras_tpc_v1
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.latest import \
+                get_keras_tpc_latest
+
+            # Keras: TPC versioning
+            tpc_models_dict = {'v1': get_keras_tpc_v1,
+                               LATEST: get_keras_tpc_latest}
+    elif fw_name == PYTORCH:
+        ###############################
+        # Build Pytorch TPC models
+        ###############################
+        if FOUND_TORCH:
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.v1.tpc_pytorch import \
+                get_pytorch_tpc as get_pytorch_tpc_v1
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.latest import \
+                get_pytorch_tpc_latest
+
+            # Pytorch: TPC versioning
+            tpc_models_dict = {'v1': get_pytorch_tpc_v1,
                                LATEST: get_pytorch_tpc_latest}
-
-tpc_dict = {TENSORFLOW: keras_tpc_models_dict,
-            PYTORCH: pytorch_tpc_models_dict}
-
-
+    if tpc_models_dict is not None:
+        return tpc_models_dict
+    else:
+        Logger.critical(f'Framework {fw_name} is not supported in imx500 or the relevant packages are not '
+                        f'installed. Please make sure the relevant packages are installed when using MCT for optimizing'
+                        f' a {fw_name} model. For Tensorflow, please install tensorflow. For PyTorch, please install '
+                        f'torch.')  # pragma: no cover

model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/target_platform_capabilities.py

@@ -14,34 +14,41 @@
 # ==============================================================================
 
 from model_compression_toolkit.constants import FOUND_TF, FOUND_TORCH, TENSORFLOW, PYTORCH
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.target_platform_capabilities.constants import LATEST
 
-
-###############################
-# Build Tensorflow TPC models
-###############################
-keras_tpc_models_dict = None
-if FOUND_TF:
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.v1.tpc_keras import get_keras_tpc as get_keras_tpc_v1
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.latest import get_keras_tpc_latest
-
-    # Keras: TPC versioning
-    keras_tpc_models_dict = {'v1': get_keras_tpc_v1,
-                             LATEST: get_keras_tpc_latest}
-
-###############################
-# Build Pytorch TPC models
-###############################
-pytorch_tpc_models_dict = None
-if FOUND_TORCH:
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.v1.tpc_pytorch import \
-        get_pytorch_tpc as get_pytorch_tpc_v1
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.latest import get_pytorch_tpc_latest
-
-    # Pytorch: TPC versioning
-    pytorch_tpc_models_dict = {'v1': get_pytorch_tpc_v1,
+def get_tpc_dict_by_fw(fw_name):
+    tpc_models_dict = None
+    if fw_name == TENSORFLOW:
+        ###############################
+        # Build Tensorflow TPC models
+        ###############################
+        if FOUND_TF:
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.v1.tpc_keras import \
+                get_keras_tpc as get_keras_tpc_v1
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.latest import \
+                get_keras_tpc_latest
+
+            # Keras: TPC versioning
+            tpc_models_dict = {'v1': get_keras_tpc_v1,
+                               LATEST: get_keras_tpc_latest}
+    elif fw_name == PYTORCH:
+        ###############################
+        # Build Pytorch TPC models
+        ###############################
+        if FOUND_TORCH:
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.v1.tpc_pytorch import \
+                get_pytorch_tpc as get_pytorch_tpc_v1
+            from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.latest import \
+                get_pytorch_tpc_latest
+
+            # Pytorch: TPC versioning
+            tpc_models_dict = {'v1': get_pytorch_tpc_v1,
                                LATEST: get_pytorch_tpc_latest}
-
-tpc_dict = {TENSORFLOW: keras_tpc_models_dict,
-            PYTORCH: pytorch_tpc_models_dict}
-
+    if tpc_models_dict is not None:
+        return tpc_models_dict
+    else:
+        Logger.critical(f'Framework {fw_name} is not supported in imx500 or the relevant packages are not '
+                        f'installed. Please make sure the relevant packages are installed when using MCT for optimizing'
+                        f' a {fw_name} model. For Tensorflow, please install tensorflow. For PyTorch, please install '
+                        f'torch.')  # pragma: no cover

model_compression_toolkit/xquant/common/constants.py

@@ -27,6 +27,8 @@ INTERMEDIATE_SIMILARITY_METRICS_VAL = 'intermediate_similarity_metrics_val'
 # Graph attribute names:
 XQUANT_REPR = 'xquant_repr'
 XQUANT_VAL = 'xquant_val'
+CUT_MEMORY_ELEMENTS = 'cut_memory_elements'
+CUT_TOTAL_SIZE = 'cut_total_size'
 
 # Report file name:
 REPORT_FILENAME = 'quant_report.json'
@@ -36,3 +38,4 @@ TENSORBOARD_DEFAULT_TAG = 'xquant'
 
 # When extracting the activations of a model we hold the output using a dedicated key:
 MODEL_OUTPUT_KEY = 'model_output_key'
+

model_compression_toolkit/xquant/common/core_report_generator.py

@@ -45,6 +45,9 @@ def core_report_generator(float_model: Any,
     Returns:
         Dict[str, Any]: A dictionary containing the collected similarity metrics and report data.
     """
+    # Get metadata from the quantized model
+    quantized_model_metadata = fw_report_utils.get_metadata_fn(quantized_model)
+
     # Collect histograms on the float model.
     float_graph = fw_report_utils.model_folding_utils.create_float_folded_graph(float_model, repr_dataset)
     mi = ModelCollector(float_graph, fw_report_utils.fw_impl, fw_report_utils.fw_info)
@@ -74,7 +77,12 @@ def core_report_generator(float_model: Any,
     # Add a graph of the quantized model with the similarity metrics to TensorBoard for visualization.
     fw_report_utils.tb_utils.add_graph_to_tensorboard(quantized_model,
                                                       similarity_metrics,
-                                                      repr_dataset)
+                                                      repr_dataset,
+                                                      quantized_model_metadata)
+
+    # Adds text information (like max cut and output similarity metrics) to the tensorboard writer.
+    fw_report_utils.tb_utils.add_text_information(similarity_metrics,
+                                                  quantized_model_metadata)
 
     # Save data to a json file.
     fw_report_utils.dump_report_to_json(report_dir=xquant_config.report_dir,

model_compression_toolkit/xquant/common/framework_report_utils.py

@@ -18,7 +18,7 @@ import os
 
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
-from typing import Any, Dict
+from typing import Any, Dict, Callable
 
 from model_compression_toolkit.xquant.common.constants import REPORT_FILENAME
 from model_compression_toolkit.xquant.common.dataset_utils import DatasetUtils
@@ -39,7 +39,8 @@ class FrameworkReportUtils:
                  similarity_calculator: SimilarityCalculator,
                  dataset_utils: DatasetUtils,
                  model_folding_utils: ModelFoldingUtils,
-                 tb_utils: TensorboardUtils):
+                 tb_utils: TensorboardUtils,
+                 get_metadata_fn: Callable):
         """
         Initializes the FrameworkReportUtils class with various utility components required for generating the report.
 
@@ -50,6 +51,7 @@ class FrameworkReportUtils:
             dataset_utils (DatasetUtils): Utilities for handling datasets.
             model_folding_utils (ModelFoldingUtils): Utilities for model folding operations.
             tb_utils (TensorboardUtils): Utilities for TensorBoard operations.
+            get_metadata_fn (Callable): Function to retrieve the metadata from the quantized model.
         """
         self.fw_info = fw_info
         self.fw_impl = fw_impl
@@ -57,18 +59,7 @@ class FrameworkReportUtils:
         self.dataset_utils = dataset_utils
         self.model_folding_utils = model_folding_utils
         self.tb_utils = tb_utils
-
-    def create_report_directory(self, dir_path: str):
-        """
-        Create a directory for saving reports.
-
-        Args:
-            dir_path (str): The path to the directory to create.
-
-        """
-        if not os.path.exists(dir_path):
-            os.makedirs(dir_path, exist_ok=True)
-            Logger.info(f"Directory created at: {dir_path}")
+        self.get_metadata_fn = get_metadata_fn
 
     def dump_report_to_json(self,
                             report_dir: str,