mct-nightly 2.1.0.20240616.65727__py3-none-any.whl → 2.1.0.20240618.432__py3-none-any.whl

model_compression_toolkit/xquant/common/tensorboard_utils.py

@@ -0,0 +1,101 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+
+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.logger import Logger
+
+
+from typing import Any, Dict, Callable
+
+
+class TensorboardUtils:
+    """
+    Utility class for handling Tensorboard operations like adding graph to display
+    and histograms on the float model.
+    """
+
+    def __init__(self,
+                 report_dir: str,
+                 fw_info: FrameworkInfo,
+                 fw_impl: FrameworkImplementation):
+        """
+        Initialize the TensorboardUtils.
+
+        Args:
+            report_dir (str): Directory where Tensorboard logs will be stored.
+            fw_info (FrameworkInfo): Framework-specific information.
+            fw_impl (FrameworkImplementation): Framework-specific implementation.
+        """
+        self.fw_impl = fw_impl
+        self.fw_info = fw_info
+        self.tb_writer = TensorboardWriter(report_dir, fw_info)
+        Logger.info(f"Please run: tensorboard --logdir {self.tb_writer.dir_path}")
+
+    def get_graph_for_tensorboard_display(self,
+                                          quantized_model: Any,
+                                          similarity_metrics: Dict[str, Any],
+                                          repr_dataset: Callable) -> Graph:
+        """
+        Get the graph for Tensorboard display. The framework-specific implementations
+        (like KerasTensorboardUtils and PytorchTensorboardUtils) should implement this
+        as it differs between them when combining the similarity metrics into the graph.
+
+        Args:
+            quantized_model (Any): The quantized model.
+            similarity_metrics (Dict[str, Any]): Metrics for model similarity.
+            repr_dataset (Callable): Representative dataset function.
+
+        Returns:
+            Graph: The generated graph for Tensorboard display.
+        """
+        Logger.critical("This method should be implemented by the framework-specific TensorboardUtils.") # pragma: no cover
+
+    def add_histograms_to_tensorboard(self,
+                                      graph: Graph):
+        """
+        Add histograms to Tensorboard from a graph that holds these statistics.
+
+        Args:
+            graph (Graph): Graph with histograms to add to the tensorboard.
+        """
+        self.tb_writer.add_histograms(graph, TENSORBOARD_DEFAULT_TAG)
+
+    def add_graph_to_tensorboard(self,
+                                 quantized_model: Any,
+                                 similarity_metrics: Dict[str, Any],
+                                 repr_dataset: Callable):
+        """
+        Add a graph to Tensorboard. The graph represents the quantized graph
+        with the similarity metrics that were measured in different nodes.
+
+        Args:
+            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).
+        """
+        # 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)
+
+        self.tb_writer.add_graph(tb_graph, TENSORBOARD_DEFAULT_TAG)
+
+

model_compression_toolkit/xquant/common/xquant_config.py

@@ -0,0 +1,39 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+
+from typing import Dict, Callable
+
+
+class XQuantConfig:
+    """
+    Configuration for generating the report.
+    It allows to set the log dir that the report will be saved in and to add similarity metrics
+    to measure between tensors of the two models.
+    """
+
+    def __init__(self,
+                 report_dir: str,
+                 custom_similarity_metrics: Dict[str, Callable] = None):
+        """
+        Initializes the configuration for explainable quantization.
+
+        Args:
+            report_dir (str): Directory where the reports will be saved.
+            custom_similarity_metrics (Dict[str, Callable]): Custom similarity metrics to be computed between tensors
+            of the two models. The dictionary keys are similarity metric names and the values are callables that implement the
+            similarity metric computation.
+        """
+        self.report_dir = report_dir
+        self.custom_similarity_metrics = custom_similarity_metrics

model_compression_toolkit/xquant/keras/__init__.py

@@ -0,0 +1,15 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+

model_compression_toolkit/xquant/keras/dataset_utils.py

@@ -0,0 +1,57 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+
+from typing import Any
+
+from model_compression_toolkit.xquant.common.dataset_utils import DatasetUtils
+
+
+class KerasDatasetUtils(DatasetUtils):
+    """
+    Class with helpful methods for handling different kinds of Keras datasets from the user.
+    """
+
+    @staticmethod
+    def prepare_dataset(dataset: Any, is_validation: bool, device: str = None):
+        """
+        Prepare the dataset so calling it will return only inputs for the model (like in the case
+        of the representative dataset). For example, when the validation dataset is used, the labels
+        should be removed.
+
+        Args:
+            dataset: Dataset to prepare.
+            is_validation: Whether it's validation dataset or not.
+            device: Device to transfer the data to.
+
+        Returns:
+            Generator to use for retrieving the dataset inputs.
+        """
+        def process_data(x: Any, is_validation: bool):
+            """
+            Processes individual data samples to transfer them to the device and convert to torch tensors if needed.
+
+            Args:
+                data: The data sample to process.
+                is_validation: A flag indicating if this is a validation dataset.
+                device: The device to transfer the data to.
+
+            Returns:
+                The data as torch tensors on the desired device.
+            """
+            return x[0] if is_validation else x  # Assume data[0] contains the inputs and data[1] the labels
+
+        for x in dataset():
+            yield process_data(x, is_validation)
+

model_compression_toolkit/xquant/keras/facade_xquant_report.py

@@ -0,0 +1,63 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+
+from typing import Callable, Dict, Any
+
+from model_compression_toolkit.constants import FOUND_TF
+from model_compression_toolkit.xquant.common.core_report_generator import core_report_generator
+from model_compression_toolkit.xquant import XQuantConfig
+from model_compression_toolkit.logger import Logger
+
+if FOUND_TF:
+    import keras
+    from model_compression_toolkit.xquant.keras.keras_report_utils import KerasReportUtils
+
+    def xquant_report_keras_experimental(float_model: keras.Model,
+                                         quantized_model: keras.Model,
+                                         repr_dataset: Callable,
+                                         validation_dataset: Callable,
+                                         xquant_config: XQuantConfig) -> Dict[str, Any]:
+        """
+        Generate an explainable quantization report for a quantized Keras model.
+
+        Args:
+            float_model (keras.Model): The original floating-point Keras model.
+            quantized_model (keras.Model): The quantized Keras model.
+            repr_dataset (Callable): The representative dataset used during quantization for similarity metrics computation.
+            validation_dataset (Callable): The validation dataset used for evaluation for similarity metrics computation.
+            xquant_config (XQuantConfig): Configuration settings for explainable quantization.
+
+        Returns:
+            Dict[str, Any]: A dictionary containing the collected similarity metrics and report data.
+        """
+        # Initialize the logger with the report directory.
+        Logger.set_log_file(log_folder=xquant_config.report_dir)
+
+        # Initialize a utility class for handling Keras-specific reporting tasks.
+        keras_report_utils = KerasReportUtils(xquant_config.report_dir)
+
+        # Create the report after collecting useful data like histograms and similarity metrics.
+        _collected_data = core_report_generator(float_model=float_model,
+                                                quantized_model=quantized_model,
+                                                repr_dataset=repr_dataset,
+                                                validation_dataset=validation_dataset,
+                                                fw_report_utils=keras_report_utils,
+                                                xquant_config=xquant_config)
+
+        return _collected_data
+else:
+    def xquant_report_keras_experimental(*args, **kwargs):
+        Logger.critical("Tensorflow must be installed to use xquant_report_keras_experimental. "
+                                     "The 'tensorflow' package is missing.")  # pragma: no cover

model_compression_toolkit/xquant/keras/keras_report_utils.py

@@ -0,0 +1,60 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+
+
+from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
+from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
+from model_compression_toolkit.xquant.common.framework_report_utils import FrameworkReportUtils
+from model_compression_toolkit.ptq.keras.quantization_facade import DEFAULT_KERAS_TPC
+from model_compression_toolkit.xquant.common.model_folding_utils import ModelFoldingUtils
+from model_compression_toolkit.xquant.common.similarity_calculator import SimilarityCalculator
+from model_compression_toolkit.xquant.keras.dataset_utils import KerasDatasetUtils
+from model_compression_toolkit.xquant.keras.model_analyzer import KerasModelAnalyzer
+
+from model_compression_toolkit.xquant.keras.similarity_functions import KerasSimilarityFunctions
+from model_compression_toolkit.xquant.keras.tensorboard_utils import KerasTensorboardUtils
+
+
+class KerasReportUtils(FrameworkReportUtils):
+    """
+    Class with various utility components required for generating the report for a Keras model.
+    """
+    def __init__(self, report_dir: str):
+        """
+        Args:
+            report_dir: Logging dir path.
+        """
+        fw_info = DEFAULT_KERAS_INFO
+        fw_impl = KerasImplementation()
+
+        dataset_utils = KerasDatasetUtils()
+        model_folding = ModelFoldingUtils(fw_info=fw_info,
+                                          fw_impl=fw_impl,
+                                          fw_default_tpc=DEFAULT_KERAS_TPC)
+
+        similarity_calculator = SimilarityCalculator(dataset_utils=dataset_utils,
+                                                     model_folding=model_folding,
+                                                     similarity_functions=KerasSimilarityFunctions(),
+                                                     model_analyzer_utils=KerasModelAnalyzer())
+
+        tb_utils = KerasTensorboardUtils(report_dir=report_dir,
+                                         fw_impl=fw_impl,
+                                         fw_info=fw_info)
+        super().__init__(fw_info,
+                         fw_impl,
+                         similarity_calculator,
+                         dataset_utils,
+                         model_folding,
+                         tb_utils)

model_compression_toolkit/xquant/keras/model_analyzer.py

@@ -0,0 +1,136 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+from typing import List, Tuple, Dict
+
+from mct_quantizers import KerasQuantizationWrapper
+from model_compression_toolkit.xquant.common.constants import MODEL_OUTPUT_KEY
+from model_compression_toolkit.xquant.common.model_analyzer import ModelAnalyzer
+import keras
+import numpy as np
+
+
+class KerasModelAnalyzer(ModelAnalyzer):
+    """
+    This class provides utilities for analyzing Keras models, specifically for
+    extracting activations and comparing float and quantized models.
+    """
+
+    def extract_model_activations(self,
+                                  float_model: keras.Model,
+                                  quantized_model: keras.Model,
+                                  float_name2quant_name: Dict[str, str],
+                                  data: List[np.ndarray]) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray]]:
+        """
+        Extracts activations from both the float and quantized models.
+
+        Args:
+            float_model (keras.Model): The float model.
+            quantized_model (keras.Model): The quantized model.
+            float_name2quant_name (Dict[str, str]): A mapping from float model layer names to quantized model layer
+            names.
+            data (List[np.ndarray]): Input data for which to compute activations.
+
+        Returns:
+            Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray]]:
+                - Dictionary of activations for the float model.
+                - Dictionary of activations for the quantized model.
+        """
+
+        def _compute_activations(model: keras.Model, layer_names: List[str], data: List[np.ndarray]) -> Dict[str, np.ndarray]:
+            """
+            Computes the activations for the specified layers of the model, and the model's output.
+
+            Args:
+                model (keras.Model): The model from which to extract activations.
+                layer_names (List[str]): Names of the layers for which to compute activations.
+                data (List[np.ndarray]): Input data for the model.
+
+            Returns:
+                Dict[str, np.ndarray]:
+                    - Dictionary mapping layer names to their corresponding activations. The model's output is stored using the key MODEL_OUTPUT_KEY.
+            """
+            # Extract the outputs of the specified layers plus the model output
+            _model_outputs = [model.get_layer(name).output for name in layer_names] + [model.output]
+
+            # Create a new model that outputs the intermediate and final layer outputs
+            intermediate_layer_model = keras.Model(inputs=model.input, outputs=_model_outputs)
+            predictions = intermediate_layer_model(data)
+
+            # Map layer names to their corresponding activations and return with the output predictions
+            activation_tensors = {layer_name: predictions[i].numpy() for i, layer_name in enumerate(layer_names)}
+            activation_tensors.update({MODEL_OUTPUT_KEY: predictions[-1].numpy()})
+            return activation_tensors
+
+        # Compute activations for the quantized model
+        quant_activations = _compute_activations(quantized_model,
+                                                 list(float_name2quant_name.values()),
+                                                 data)
+        # Compute activations for the float model
+        float_activations = _compute_activations(float_model,
+                                                 list(float_name2quant_name.keys()),
+                                                 data)
+
+        # Concatenate predictions if they are lists.
+        if isinstance(quant_activations[MODEL_OUTPUT_KEY], list):
+            quant_activations[MODEL_OUTPUT_KEY] = np.concatenate(quant_activations[MODEL_OUTPUT_KEY])
+        if isinstance(float_activations[MODEL_OUTPUT_KEY], list):
+            float_activations[MODEL_OUTPUT_KEY] = np.concatenate(float_activations[MODEL_OUTPUT_KEY])
+
+        return float_activations, quant_activations
+
+    def identify_quantized_compare_points(self, quantized_model: keras.Model) -> List[str]:
+        """
+        Identifies the layers in the quantized model that are wrapped with the quantization wrapper.
+        These layers will serve as comparison points.
+
+        Notes:
+            This currently means that the quantized compare points are the linear layers that are wrapped,
+            but this may be changed in the future.
+
+        Args:
+            quantized_model (keras.Model): The quantized model from which to identify comparison points.
+
+        Returns:
+            List[str]: Names of the layers wrapped with the quantization wrapper.
+        """
+        return [layer.name for layer in quantized_model.layers if isinstance(layer, KerasQuantizationWrapper)]
+
+    def find_corresponding_float_layer(self,
+                                       quant_compare_point: str,
+                                       quantized_model: keras.Model) -> str:
+        """
+        Finds the corresponding float model layer for a given quantized model layer.
+
+        Args:
+            quant_compare_point (str): The name of the quantized model layer.
+            quantized_model (keras.Model): The quantized model.
+
+        Returns:
+            str: The name of the corresponding layer in the float model.
+        """
+        return quantized_model.get_layer(quant_compare_point).layer.name
+
+    def extract_float_layer_names(self, float_model: keras.Model) -> List[str]:
+        """
+        Extracts the names of all layers in the float model.
+
+        Args:
+            float_model (keras.Model): The float model from which to extract layer names.
+
+        Returns:
+            List[str]: Names of all layers in the float model.
+        """
+        float_layers_names = [layer.name for layer in float_model.layers]
+        return float_layers_names

model_compression_toolkit/xquant/keras/similarity_functions.py

@@ -0,0 +1,75 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+
+import numpy as np
+import tensorflow as tf
+
+from model_compression_toolkit.xquant.common.similarity_functions import SimilarityFunctions
+
+class KerasSimilarityFunctions(SimilarityFunctions):
+    """
+    A class that extends SimilarityFunctions to implement similarity metrics using Keras.
+    Even though the names referred to are quantized and float, it can help compare between
+    tensors of any two models.
+    """
+
+    @staticmethod
+    def compute_mse(x: np.ndarray, y: np.ndarray) -> float:
+        """
+        Compute the Mean Squared Error (MSE) between two tensors (usually, the float and quantized tensors).
+
+        Args:
+            x (np.ndarray): First tensor to compare.
+            y (np.ndarray): Second tensor to compare.
+
+        Returns:
+            float: The computed MSE value.
+        """
+        mse = tf.keras.losses.MeanSquaredError()(x, y)
+        return float(mse.numpy())
+
+    @staticmethod
+    def compute_cs(x: np.ndarray, y: np.ndarray) -> float:
+        """
+        Compute the Cosine Similarity (CS) between two tensors (usually, the float and quantized tensors).
+
+        Args:
+            x (np.ndarray): First tensor to compare.
+            y (np.ndarray): Second tensor to compare.
+
+        Returns:
+            float: The computed CS value.
+        """
+        cs = tf.keras.losses.CosineSimilarity()(x.flatten(), y.flatten())
+        return float(cs.numpy())
+
+    @staticmethod
+    def compute_sqnr(x: np.ndarray, y: np.ndarray) -> float:
+        """
+        Compute the Signal-to-Quantization-Noise Ratio (SQNR) between two tensors (usually, the float and quantized tensors).
+
+        Args:
+            x (np.ndarray): First tensor to compare.
+            y (np.ndarray): Second tensor to compare.
+
+        Returns:
+            float: The computed SQNR value.
+        """
+        signal_power = tf.reduce_mean(tf.square(x))
+        noise_power = tf.reduce_mean(tf.square(x - y))
+        sqnr = signal_power / noise_power
+        return float(sqnr.numpy())
+
+

model_compression_toolkit/xquant/keras/tensorboard_utils.py

@@ -0,0 +1,84 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+from typing import Dict, Callable
+
+import keras
+
+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.keras.reader.reader import model_reader
+
+from model_compression_toolkit.xquant.common.constants import XQUANT_REPR, INTERMEDIATE_SIMILARITY_METRICS_REPR, XQUANT_VAL, INTERMEDIATE_SIMILARITY_METRICS_VAL
+from model_compression_toolkit.xquant.common.tensorboard_utils import TensorboardUtils
+
+
+class KerasTensorboardUtils(TensorboardUtils):
+    """
+    A utility class for handling TensorBoard operations specific to Keras models.
+    This class extends the generic TensorboardUtils class and provides methods
+    to facilitate the visualization of quantized models and their similarity metrics
+    in TensorBoard.
+    """
+
+    def __init__(self, report_dir: str,
+                 fw_info: FrameworkInfo,
+                 fw_impl: FrameworkImplementation):
+        """
+        Initialize the KerasTensorboardUtils class with the given parameters.
+
+        Args:
+            report_dir (str): Directory where the TensorBoard files will be stored.
+            fw_info (FrameworkInfo): Information about the framework being used.
+            fw_impl (FrameworkImplementation): Implementation functions for the framework.
+        """
+        super().__init__(report_dir,
+                         fw_info,
+                         fw_impl)
+
+    def get_graph_for_tensorboard_display(self,
+                                          quantized_model: keras.Model,
+                                          similarity_metrics: Dict[str, Dict[str, float]],
+                                          repr_dataset: Callable) -> Graph:
+        """
+        Generate a graph suitable for TensorBoard display from the provided quantized model
+        and similarity metrics.
+
+        Args:
+            quantized_model (keras.Model): The quantized Keras model for which the graph is to be created.
+            similarity_metrics (Dict[str, Dict[str, float]]): A dictionary containing similarity metrics
+                for different nodes in the model.
+            repr_dataset (Callable): A function or callable that provides the representative dataset.
+
+        Returns:
+            Graph: A graph object representing the quantized model, annotated with similarity metrics.
+        """
+        # Read the quantized model into a graph structure.
+        quant_graph = model_reader(quantized_model)
+
+        # Iterate over each node in the graph.
+        for node in quant_graph.nodes:
+            # Check if the node's name is in the similarity metrics for intermediate representation.
+            if node.name in similarity_metrics[INTERMEDIATE_SIMILARITY_METRICS_REPR].keys():
+                # If so, add the similarity metric for intermediate representation to the node's attributes.
+                node.framework_attr[XQUANT_REPR] = similarity_metrics[INTERMEDIATE_SIMILARITY_METRICS_REPR][node.name]
+
+            # Check if the node's name is in the similarity metrics for validation.
+            if node.name in similarity_metrics[INTERMEDIATE_SIMILARITY_METRICS_VAL].keys():
+                # If so, add the similarity metric for validation to the node's attributes.
+                node.framework_attr[XQUANT_VAL] = similarity_metrics[INTERMEDIATE_SIMILARITY_METRICS_VAL][node.name]
+
+        return quant_graph

model_compression_toolkit/xquant/pytorch/__init__.py

@@ -0,0 +1,15 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      http://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+#