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

model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_pot/tp_model.py

@@ -18,7 +18,7 @@ import model_compression_toolkit as mct
 from model_compression_toolkit.constants import FLOAT_BITWIDTH
 from model_compression_toolkit.target_platform_capabilities.constants import KERNEL_ATTR, BIAS_ATTR, WEIGHTS_N_BITS
 from model_compression_toolkit.target_platform_capabilities.target_platform import OpQuantizationConfig, \
-    TargetPlatformModel
+    TargetPlatformModel, Signedness
 from model_compression_toolkit.target_platform_capabilities.target_platform.op_quantization_config import \
     AttributeQuantizationConfig
 
@@ -94,7 +94,8 @@ def get_op_quantization_configs() -> Tuple[OpQuantizationConfig, List[OpQuantiza
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=32)
+        simd_size=32,
+        signedness=Signedness.AUTO)
 
     # We define an 8-bit config for linear operations quantization, that include a kernel and bias attributes.
     linear_eight_bits = tp.OpQuantizationConfig(
@@ -107,7 +108,8 @@ def get_op_quantization_configs() -> Tuple[OpQuantizationConfig, List[OpQuantiza
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=32)
+        simd_size=32,
+        signedness=Signedness.AUTO)
 
     # To quantize a model using mixed-precision, create
     # a list with more than one OpQuantizationConfig.

model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2/tp_model.py

@@ -18,7 +18,7 @@ import model_compression_toolkit as mct
 from model_compression_toolkit.constants import FLOAT_BITWIDTH
 from model_compression_toolkit.target_platform_capabilities.constants import KERNEL_ATTR, BIAS_ATTR, WEIGHTS_N_BITS
 from model_compression_toolkit.target_platform_capabilities.target_platform import OpQuantizationConfig, \
-    TargetPlatformModel
+    TargetPlatformModel, Signedness
 from model_compression_toolkit.target_platform_capabilities.target_platform.op_quantization_config import \
     AttributeQuantizationConfig
 
@@ -100,7 +100,8 @@ def get_op_quantization_configs() -> \
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=32)
+        simd_size=32,
+        signedness=Signedness.AUTO)
 
     # We define an 8-bit config for linear operations quantization, that include a kernel and bias attributes.
     linear_eight_bits = tp.OpQuantizationConfig(
@@ -113,7 +114,8 @@ def get_op_quantization_configs() -> \
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=32)
+        simd_size=32,
+        signedness=Signedness.AUTO)
 
     # To quantize a model using mixed-precision, create
     # a list with more than one OpQuantizationConfig.

model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_lut/tp_model.py

@@ -19,7 +19,7 @@ from model_compression_toolkit.constants import FLOAT_BITWIDTH
 from model_compression_toolkit.target_platform_capabilities.constants import KERNEL_ATTR, BIAS_ATTR, WEIGHTS_N_BITS, \
     WEIGHTS_QUANTIZATION_METHOD
 from model_compression_toolkit.target_platform_capabilities.target_platform import OpQuantizationConfig, \
-    TargetPlatformModel
+    TargetPlatformModel, Signedness
 from model_compression_toolkit.target_platform_capabilities.target_platform.op_quantization_config import \
     AttributeQuantizationConfig
 
@@ -96,7 +96,8 @@ def get_op_quantization_configs() -> \
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=32)
+        simd_size=32,
+        signedness=Signedness.AUTO)
 
     # We define an 8-bit config for linear operations quantization, that include a kernel and bias attributes.
     linear_eight_bits = tp.OpQuantizationConfig(
@@ -109,7 +110,8 @@ def get_op_quantization_configs() -> \
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=32)
+        simd_size=32,
+        signedness=Signedness.AUTO)
 
     # To quantize a model using mixed-precision, create
     # a list with more than one OpQuantizationConfig.

model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3/tp_model.py

@@ -18,7 +18,7 @@ import model_compression_toolkit as mct
 from model_compression_toolkit.constants import FLOAT_BITWIDTH
 from model_compression_toolkit.target_platform_capabilities.constants import KERNEL_ATTR, BIAS_ATTR, WEIGHTS_N_BITS
 from model_compression_toolkit.target_platform_capabilities.target_platform import OpQuantizationConfig, \
-    TargetPlatformModel
+    TargetPlatformModel, Signedness
 from model_compression_toolkit.target_platform_capabilities.target_platform.op_quantization_config import \
     AttributeQuantizationConfig
 
@@ -100,7 +100,8 @@ def get_op_quantization_configs() -> \
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=32)
+        simd_size=32,
+        signedness=Signedness.AUTO)
 
     # We define an 8-bit config for linear operations quantization, that include a kernel and bias attributes.
     linear_eight_bits = tp.OpQuantizationConfig(
@@ -113,7 +114,8 @@ def get_op_quantization_configs() -> \
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=32)
+        simd_size=32,
+        signedness=Signedness.AUTO)
 
     # To quantize a model using mixed-precision, create
     # a list with more than one OpQuantizationConfig.

model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_lut/tp_model.py

@@ -19,7 +19,7 @@ from model_compression_toolkit.constants import FLOAT_BITWIDTH
 from model_compression_toolkit.target_platform_capabilities.constants import KERNEL_ATTR, BIAS_ATTR, WEIGHTS_N_BITS, \
     WEIGHTS_QUANTIZATION_METHOD
 from model_compression_toolkit.target_platform_capabilities.target_platform import OpQuantizationConfig, \
-    TargetPlatformModel
+    TargetPlatformModel, Signedness
 from model_compression_toolkit.target_platform_capabilities.target_platform.op_quantization_config import \
     AttributeQuantizationConfig
 
@@ -96,7 +96,8 @@ def get_op_quantization_configs() -> \
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=32)
+        simd_size=32,
+        signedness=Signedness.AUTO)
 
     # We define an 8-bit config for linear operations quantization, that include a kernel and bias attributes.
     linear_eight_bits = tp.OpQuantizationConfig(
@@ -109,7 +110,8 @@ def get_op_quantization_configs() -> \
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=32)
+        simd_size=32,
+        signedness=Signedness.AUTO)
 
     # To quantize a model using mixed-precision, create
     # a list with more than one OpQuantizationConfig.

model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tp_model.py

@@ -18,7 +18,7 @@ import model_compression_toolkit as mct
 from model_compression_toolkit.constants import FLOAT_BITWIDTH
 from model_compression_toolkit.target_platform_capabilities.constants import KERNEL_ATTR, BIAS_ATTR, WEIGHTS_N_BITS
 from model_compression_toolkit.target_platform_capabilities.target_platform import OpQuantizationConfig, \
-    TargetPlatformModel
+    TargetPlatformModel, Signedness
 from model_compression_toolkit.target_platform_capabilities.target_platform.op_quantization_config import \
     AttributeQuantizationConfig
 
@@ -100,7 +100,8 @@ def get_op_quantization_configs() -> \
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=32)
+        simd_size=32,
+        signedness=Signedness.AUTO)
 
     # We define an 8-bit config for linear operations quantization, that include a kernel and bias attributes.
     linear_eight_bits = tp.OpQuantizationConfig(
@@ -113,7 +114,8 @@ def get_op_quantization_configs() -> \
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=32)
+        simd_size=32,
+        signedness=Signedness.AUTO)
 
     # To quantize a model using mixed-precision, create
     # a list with more than one OpQuantizationConfig.
@@ -170,7 +172,7 @@ def generate_tp_model(default_config: OpQuantizationConfig,
     const_config_input16 = const_config.clone_and_edit(
         supported_input_activation_n_bits=(8, 16))
     const_config_input16_output16 = const_config_input16.clone_and_edit(
-        activation_n_bits=16, is_signed=True)
+        activation_n_bits=16, signedness=Signedness.SIGNED)
     const_configuration_options_inout16 = tp.QuantizationConfigOptions([const_config_input16_output16,
                                                                         const_config_input16],
                                                                        base_config=const_config_input16)

model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_pytorch.py

@@ -17,7 +17,7 @@ import operator
 
 import torch
 from torch import add, sub, mul, div, flatten, reshape, split, unsqueeze, dropout, sigmoid, tanh, chunk, unbind, topk, \
-    gather, equal, transpose, permute, argmax, squeeze
+    gather, equal, transpose, permute, argmax, squeeze, multiply, subtract
 from torch.nn import Conv2d, Linear, ConvTranspose2d, MaxPool2d
 from torch.nn import Dropout, Flatten, Hardtanh, Identity
 from torch.nn import ReLU, ReLU6, PReLU, SiLU, Sigmoid, Tanh, Hardswish, LeakyReLU
@@ -101,8 +101,8 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                              tp.LayerFilterParams(hardtanh, min_val=0)])
 
         tp.OperationsSetToLayers("Add", [operator.add, add])
-        tp.OperationsSetToLayers("Sub", [operator.sub, sub])
-        tp.OperationsSetToLayers("Mul", [operator.mul, mul])
+        tp.OperationsSetToLayers("Sub", [operator.sub, sub, subtract])
+        tp.OperationsSetToLayers("Mul", [operator.mul, mul, multiply])
         tp.OperationsSetToLayers("Div", [operator.truediv, div])
         tp.OperationsSetToLayers("PReLU", [PReLU, prelu])
         tp.OperationsSetToLayers("Swish", [SiLU, silu, Hardswish, hardswish])

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/qnnpack_tpc/v1/tp_model.py

@@ -18,7 +18,7 @@ import model_compression_toolkit as mct
 from model_compression_toolkit.constants import FLOAT_BITWIDTH
 from model_compression_toolkit.target_platform_capabilities.constants import KERNEL_ATTR, BIAS_ATTR
 from model_compression_toolkit.target_platform_capabilities.target_platform import OpQuantizationConfig, \
-    TargetPlatformModel
+    TargetPlatformModel, Signedness
 from model_compression_toolkit.target_platform_capabilities.target_platform.op_quantization_config import \
     AttributeQuantizationConfig
 
@@ -95,7 +95,8 @@ def get_op_quantization_configs() -> Tuple[OpQuantizationConfig, List[OpQuantiza
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=32)
+        simd_size=32,
+        signedness=Signedness.AUTO)
 
     # We define an 8-bit config for linear operations quantization, that include a kernel and bias attributes.
     linear_eight_bits = tp.OpQuantizationConfig(
@@ -108,8 +109,8 @@ def get_op_quantization_configs() -> Tuple[OpQuantizationConfig, List[OpQuantiza
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=None
-    )
+        simd_size=None,
+        signedness=Signedness.AUTO)
 
     mixed_precision_cfg_list = []  # No mixed precision
 

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/target_platform_capabilities/tpc_models/tflite_tpc/v1/tp_model.py

@@ -18,7 +18,7 @@ import model_compression_toolkit as mct
 from model_compression_toolkit.constants import FLOAT_BITWIDTH
 from model_compression_toolkit.target_platform_capabilities.constants import BIAS_ATTR, KERNEL_ATTR
 from model_compression_toolkit.target_platform_capabilities.target_platform import OpQuantizationConfig, \
-    TargetPlatformModel
+    TargetPlatformModel, Signedness
 from model_compression_toolkit.target_platform_capabilities.target_platform.op_quantization_config import \
     QuantizationMethod, AttributeQuantizationConfig
 
@@ -93,7 +93,8 @@ def get_op_quantization_configs() -> Tuple[OpQuantizationConfig, List[OpQuantiza
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=32)
+        simd_size=32,
+        signedness=Signedness.AUTO)
 
     # We define an 8-bit config for linear operations quantization, that include a kernel and bias attributes.
     linear_eight_bits = tp.OpQuantizationConfig(
@@ -106,8 +107,8 @@ def get_op_quantization_configs() -> Tuple[OpQuantizationConfig, List[OpQuantiza
         quantization_preserving=False,
         fixed_scale=None,
         fixed_zero_point=None,
-        simd_size=None
-    )
+        simd_size=None,
+        signedness=Signedness.AUTO)
 
     mixed_precision_cfg_list = []  # No mixed precision
 

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,

model_compression_toolkit/xquant/common/tensorboard_utils.py

@@ -12,18 +12,20 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #  ==============================================================================
-
+from model_compression_toolkit.constants import MAX_CUT
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 
 
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import TensorboardWriter
-from model_compression_toolkit.xquant.common.constants import TENSORBOARD_DEFAULT_TAG
+from model_compression_toolkit.xquant.common.constants import TENSORBOARD_DEFAULT_TAG, OUTPUT_SIMILARITY_METRICS_REPR, \
+    OUTPUT_SIMILARITY_METRICS_VAL
 from model_compression_toolkit.logger import Logger
 
 
 from typing import Any, Dict, Callable
+from mct_quantizers.keras.metadata import get_metadata
 
 
 class TensorboardUtils:
@@ -52,7 +54,8 @@ class TensorboardUtils:
     def get_graph_for_tensorboard_display(self,
                                           quantized_model: Any,
                                           similarity_metrics: Dict[str, Any],
-                                          repr_dataset: Callable) -> Graph:
+                                          repr_dataset: Callable,
+                                          quantized_model_metadata: Dict) -> Graph:
         """
         Get the graph for Tensorboard display. The framework-specific implementations
         (like KerasTensorboardUtils and PytorchTensorboardUtils) should implement this
@@ -62,6 +65,7 @@ class TensorboardUtils:
             quantized_model (Any): The quantized model.
             similarity_metrics (Dict[str, Any]): Metrics for model similarity.
             repr_dataset (Callable): Representative dataset function.
+            quantized_model_metadata (Dict): Metadata from the quantized model.
 
         Returns:
             Graph: The generated graph for Tensorboard display.
@@ -81,7 +85,8 @@ class TensorboardUtils:
     def add_graph_to_tensorboard(self,
                                  quantized_model: Any,
                                  similarity_metrics: Dict[str, Any],
-                                 repr_dataset: Callable):
+                                 repr_dataset: Callable,
+                                 quantized_model_metadata: Dict):
         """
         Add a graph to Tensorboard. The graph represents the quantized graph
         with the similarity metrics that were measured in different nodes.
@@ -90,12 +95,32 @@ class TensorboardUtils:
             quantized_model (Any): The quantized model.
             similarity_metrics (Dict[str, Any]): The similarity metrics that were collected.
             repr_dataset (Callable): Representative dataset to use (if needed, like in pytorch case).
+            quantized_model_metadata (Dict): Metadata from the quantized model.
         """
         # Generate the quantized graph with similarity metrics.
         tb_graph = self.get_graph_for_tensorboard_display(quantized_model=quantized_model,
                                                           similarity_metrics=similarity_metrics,
-                                                          repr_dataset=repr_dataset)
+                                                          repr_dataset=repr_dataset,
+                                                          quantized_model_metadata=quantized_model_metadata)
 
         self.tb_writer.add_graph(tb_graph, TENSORBOARD_DEFAULT_TAG)
 
+    def add_text_information(self,
+                             similarity_metrics: Dict[str, Dict[str, float]],
+                             quantized_model_metadata: Dict[str, Any]):
+        """
+        Adds text information (like max cut and output similarity metrics) to the tensorboard writer.
 
+        Args:
+            similarity_metrics (Dict[str, Dict[str, float]]): A dictionary containing similarity metrics between quantized and float models for both representative and validation datasets.
+            quantized_model_metadata (Dict): Metadata from the quantized model.
+        """
+        # Add the computed max cut
+        maxcut_str = f"MaxCut: {quantized_model_metadata['scheduling_info'][MAX_CUT]}"
+        self.tb_writer.add_text(maxcut_str, MAX_CUT)
+
+        # Add output similarity between quantized and float models on representative and validation datasets
+        output_similarity_repr = f"Similarity Metrics on outputs using representative dataset: \n" + "\n".join([f"{key}: {value:.4f}" for key, value in similarity_metrics[OUTPUT_SIMILARITY_METRICS_REPR].items()])
+        output_similarity_val = f"Similarity Metrics on outputs using validation dataset: \n" + "\n".join([f"{key}: {value:.4f}" for key, value in similarity_metrics[OUTPUT_SIMILARITY_METRICS_VAL].items()])
+        self.tb_writer.add_text(output_similarity_repr, OUTPUT_SIMILARITY_METRICS_REPR)
+        self.tb_writer.add_text(output_similarity_val, OUTPUT_SIMILARITY_METRICS_VAL)

model_compression_toolkit/xquant/keras/facade_xquant_report.py

@@ -56,6 +56,8 @@ if FOUND_TF:
                                                 fw_report_utils=keras_report_utils,
                                                 xquant_config=xquant_config)
 
+        Logger.shutdown()
+
         return _collected_data
 else:
     def xquant_report_keras_experimental(*args, **kwargs):

model_compression_toolkit/xquant/keras/keras_report_utils.py

@@ -25,6 +25,7 @@ from model_compression_toolkit.xquant.keras.model_analyzer import KerasModelAnal
 
 from model_compression_toolkit.xquant.keras.similarity_functions import KerasSimilarityFunctions
 from model_compression_toolkit.xquant.keras.tensorboard_utils import KerasTensorboardUtils
+from mct_quantizers.keras.metadata import get_metadata
 
 
 class KerasReportUtils(FrameworkReportUtils):
@@ -57,4 +58,5 @@ class KerasReportUtils(FrameworkReportUtils):
                          similarity_calculator,
                          dataset_utils,
                          model_folding,
-                         tb_utils)
+                         tb_utils,
+                         get_metadata)