mct-nightly 2.1.0.20240617.451__tar.gz → 2.1.0.20240618.432__tar.gz

{mct-nightly-2.1.0.20240617.451 → mct-nightly-2.1.0.20240618.432}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.1.0.20240617.451
+Version: 2.1.0.20240618.432
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct-nightly-2.1.0.20240617.451 → mct-nightly-2.1.0.20240618.432}/mct_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.1.0.20240617.451
+Version: 2.1.0.20240618.432
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct-nightly-2.1.0.20240617.451 → mct-nightly-2.1.0.20240618.432}/mct_nightly.egg-info/SOURCES.txt

@@ -499,4 +499,30 @@ model_compression_toolkit/trainable_infrastructure/keras/load_model.py
 model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py
 model_compression_toolkit/trainable_infrastructure/keras/quantizer_utils.py
 model_compression_toolkit/trainable_infrastructure/pytorch/__init__.py
-model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py
+model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py
+model_compression_toolkit/xquant/__init__.py
+model_compression_toolkit/xquant/common/__init__.py
+model_compression_toolkit/xquant/common/constants.py
+model_compression_toolkit/xquant/common/core_report_generator.py
+model_compression_toolkit/xquant/common/dataset_utils.py
+model_compression_toolkit/xquant/common/framework_report_utils.py
+model_compression_toolkit/xquant/common/model_analyzer.py
+model_compression_toolkit/xquant/common/model_folding_utils.py
+model_compression_toolkit/xquant/common/similarity_calculator.py
+model_compression_toolkit/xquant/common/similarity_functions.py
+model_compression_toolkit/xquant/common/tensorboard_utils.py
+model_compression_toolkit/xquant/common/xquant_config.py
+model_compression_toolkit/xquant/keras/__init__.py
+model_compression_toolkit/xquant/keras/dataset_utils.py
+model_compression_toolkit/xquant/keras/facade_xquant_report.py
+model_compression_toolkit/xquant/keras/keras_report_utils.py
+model_compression_toolkit/xquant/keras/model_analyzer.py
+model_compression_toolkit/xquant/keras/similarity_functions.py
+model_compression_toolkit/xquant/keras/tensorboard_utils.py
+model_compression_toolkit/xquant/pytorch/__init__.py
+model_compression_toolkit/xquant/pytorch/dataset_utils.py
+model_compression_toolkit/xquant/pytorch/facade_xquant_report.py
+model_compression_toolkit/xquant/pytorch/model_analyzer.py
+model_compression_toolkit/xquant/pytorch/pytorch_report_utils.py
+model_compression_toolkit/xquant/pytorch/similarity_functions.py
+model_compression_toolkit/xquant/pytorch/tensorboard_utils.py

{mct-nightly-2.1.0.20240617.451 → mct-nightly-2.1.0.20240618.432}/mct_nightly.egg-info/requires.txt

@@ -1,7 +1,7 @@
 networkx!=2.8.1
 tqdm
 Pillow
-numpy
+numpy<2.0
 opencv-python
 scikit-image
 scikit-learn

{mct-nightly-2.1.0.20240617.451 → mct-nightly-2.1.0.20240618.432}/model_compression_toolkit/__init__.py

@@ -27,4 +27,4 @@ from model_compression_toolkit import data_generation
 from model_compression_toolkit import pruning
 from model_compression_toolkit.trainable_infrastructure.keras.load_model import keras_load_quantized_model
 
-__version__ = "2.1.0.20240617.000451"
+__version__ = "2.1.0.20240618.000432"

{mct-nightly-2.1.0.20240617.451 → mct-nightly-2.1.0.20240618.432}/model_compression_toolkit/gptq/keras/gptq_training.py

@@ -353,7 +353,7 @@ class KerasGPTQTrainer(GPTQTrainer):
                     node.final_activation_quantization_cfg.set_quant_config_attr(config_attr, config_value)
                 if self.gptq_config.train_bias:
                     use_bias = layer.layer.get_config().get(USE_BIAS)
-                    if use_bias is not None and use_bias:
+                    if use_bias is not None and use_bias and layer.layer.bias is not None:
                         new_bias = layer.layer.bias.numpy()
                         node.set_weights_by_keys(BIAS, new_bias)
 

{mct-nightly-2.1.0.20240617.451 → mct-nightly-2.1.0.20240618.432}/model_compression_toolkit/gptq/keras/graph_info.py

@@ -63,7 +63,7 @@ def get_gptq_trainable_parameters(fxp_model: Model,
                 kernel_ops_attrs = fw_info.kernel_ops_attributes_mapping.get(type(layer.layer))
                 use_bias = kernel_ops_attrs is not None and kernel_ops_attrs[0] is not None \
                            and layer.layer.get_config().get(USE_BIAS)
-                if use_bias is not None and use_bias:
+                if use_bias is not None and use_bias and layer.layer.bias is not None:
                     bias_weights.append([layer.layer.bias])
 
     return trainable_weights, bias_weights, trainable_threshold

{mct-nightly-2.1.0.20240617.451 → mct-nightly-2.1.0.20240618.432}/model_compression_toolkit/gptq/pytorch/gptq_training.py

@@ -299,7 +299,9 @@ class PytorchGPTQTrainer(GPTQTrainer):
                 for config_attr, config_value in activation_quant_config.items():
                     node.final_activation_quantization_cfg.set_quant_config_attr(config_attr, config_value)
                 if self.gptq_config.train_bias and hasattr(layer.layer, BIAS):
-                    node.set_weights_by_keys(BIAS, self.fw_impl.to_numpy(getattr(layer.layer, BIAS)))
+                    bias = getattr(layer.layer, BIAS)
+                    if bias is not None:
+                        node.set_weights_by_keys(BIAS, self.fw_impl.to_numpy(bias))
 
         return graph_quant
 
@@ -316,4 +318,5 @@ class PytorchGPTQTrainer(GPTQTrainer):
             if isinstance(layer, PytorchQuantizationWrapper):
                 if hasattr(layer.layer, BIAS):
                     bias = getattr(layer.layer, BIAS)
-                    bias.requires_grad = self.gptq_config.train_bias
+                    if bias is not None:
+                        bias.requires_grad = self.gptq_config.train_bias

{mct-nightly-2.1.0.20240617.451 → mct-nightly-2.1.0.20240618.432}/model_compression_toolkit/gptq/pytorch/graph_info.py

@@ -56,7 +56,8 @@ def get_gptq_trainable_parameters(fxp_model: nn.Module,
 
             if add_bias and hasattr(layer.layer, BIAS):
                 bias = getattr(layer.layer, BIAS)
-                trainable_bias.append(bias)
+                if bias is not None:
+                    trainable_bias.append(bias)
 
     return trainable_aux_weights, trainable_bias, trainable_threshold
 

{mct-nightly-2.1.0.20240617.451 → mct-nightly-2.1.0.20240618.432}/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3/tpc_pytorch.py

@@ -18,7 +18,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
-from torch.nn import Conv2d, Linear, ConvTranspose2d
+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
 from torch.nn.functional import relu, relu6, prelu, silu, hardtanh, hardswish, leaky_relu
@@ -83,7 +83,8 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                                     argmax,
                                                     gather,
                                                     topk,
-                                                    squeeze])
+                                                    squeeze,
+                                                    MaxPool2d])
 
         tp.OperationsSetToLayers("Conv", [Conv2d, ConvTranspose2d],
                                  attr_mapping=pytorch_linear_attr_mapping)

{mct-nightly-2.1.0.20240617.451 → mct-nightly-2.1.0.20240618.432}/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_lut/tpc_pytorch.py

@@ -18,7 +18,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
-from torch.nn import Conv2d, Linear, ConvTranspose2d
+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
 from torch.nn.functional import relu, relu6, prelu, silu, hardtanh, hardswish, leaky_relu
@@ -82,7 +82,8 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                                     argmax,
                                                     gather,
                                                     topk,
-                                                    squeeze])
+                                                    squeeze,
+                                                    MaxPool2d])
 
         tp.OperationsSetToLayers("Conv", [Conv2d, ConvTranspose2d],
                                  attr_mapping=pytorch_linear_attr_mapping)

mct-nightly-2.1.0.20240618.432/model_compression_toolkit/xquant/__init__.py

@@ -0,0 +1,19 @@
+#  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.xquant.common.xquant_config import XQuantConfig
+from model_compression_toolkit.xquant.keras.facade_xquant_report import xquant_report_keras_experimental
+from model_compression_toolkit.xquant.pytorch.facade_xquant_report import xquant_report_pytorch_experimental
+

mct-nightly-2.1.0.20240618.432/model_compression_toolkit/xquant/common/__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.
+#  ==============================================================================
+#

mct-nightly-2.1.0.20240618.432/model_compression_toolkit/xquant/common/constants.py

@@ -0,0 +1,38 @@
+#  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.
+#  ==============================================================================
+
+# Default similarity metric names:
+CS_SIMILARITY_METRIC_NAME = 'cs'
+SQNR_SIMILARITY_METRIC_NAME = 'sqnr'
+MSE_SIMILARITY_METRIC_NAME = 'mse'
+
+# Report components names:
+OUTPUT_SIMILARITY_METRICS_REPR = 'output_similarity_metrics_repr'
+OUTPUT_SIMILARITY_METRICS_VAL = 'output_similarity_metrics_val'
+INTERMEDIATE_SIMILARITY_METRICS_REPR = 'intermediate_similarity_metrics_repr'
+INTERMEDIATE_SIMILARITY_METRICS_VAL = 'intermediate_similarity_metrics_val'
+
+# Graph attribute names:
+XQUANT_REPR = 'xquant_repr'
+XQUANT_VAL = 'xquant_val'
+
+# Report file name:
+REPORT_FILENAME = 'quant_report.json'
+
+# Tag to use in tensorboard for the graph we plot:
+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'

mct-nightly-2.1.0.20240618.432/model_compression_toolkit/xquant/common/core_report_generator.py

@@ -0,0 +1,83 @@
+#  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 tqdm import tqdm
+from typing import Callable, Any, Dict
+
+from model_compression_toolkit.core.common.model_collector import ModelCollector
+from model_compression_toolkit.xquant import XQuantConfig
+from model_compression_toolkit.xquant.common.constants import OUTPUT_SIMILARITY_METRICS_REPR, OUTPUT_SIMILARITY_METRICS_VAL, INTERMEDIATE_SIMILARITY_METRICS_REPR, \
+    INTERMEDIATE_SIMILARITY_METRICS_VAL
+from model_compression_toolkit.xquant.common.framework_report_utils import FrameworkReportUtils
+
+
+def core_report_generator(float_model: Any,
+                          quantized_model: Any,
+                          repr_dataset: Callable,
+                          validation_dataset: Callable,
+                          fw_report_utils: FrameworkReportUtils,
+                          xquant_config: XQuantConfig) -> Dict[str, Any]:
+    """
+    Generate report in tensorboard with a graph of the quantized model and similarity metrics that
+    have been measured when comparing to the float model (or any other two models).
+    The report also contains histograms that are collected on the baseline model (usually, the float
+    model).
+
+    Args:
+        float_model (Any): The original floating-point model.
+        quantized_model (Any): The model after quantization.
+        repr_dataset (Callable): Representative dataset used for similarity metrics computation.
+        validation_dataset (Callable): Validation dataset used for similarity metrics computation.
+        fw_report_utils (FrameworkReportUtils): Utilities for generating framework-specific reports.
+        xquant_config (XQuantConfig): Configuration settings for explainable quantization.
+
+    Returns:
+        Dict[str, Any]: A dictionary containing the collected similarity metrics and report data.
+    """
+    # 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)
+    for _data in tqdm(repr_dataset(), desc="Collecting Histograms"):
+        mi.infer(_data)
+
+    # Collect histograms and add them to Tensorboard.
+    fw_report_utils.tb_utils.add_histograms_to_tensorboard(graph=float_graph)
+
+    # Compute similarity metrics on representative dataset and validation set.
+    repr_similarity = fw_report_utils.similarity_calculator.compute_similarity_metrics(float_model=float_model,
+                                                                                       quantized_model=quantized_model,
+                                                                                       dataset=repr_dataset,
+                                                                                       custom_similarity_metrics=xquant_config.custom_similarity_metrics)
+    val_similarity = fw_report_utils.similarity_calculator.compute_similarity_metrics(float_model=float_model,
+                                                                                      quantized_model=quantized_model,
+                                                                                      dataset=validation_dataset,
+                                                                                      custom_similarity_metrics=xquant_config.custom_similarity_metrics,
+                                                                                      is_validation=True)
+    similarity_metrics = {
+        OUTPUT_SIMILARITY_METRICS_REPR: repr_similarity[0],
+        OUTPUT_SIMILARITY_METRICS_VAL: val_similarity[0],
+        INTERMEDIATE_SIMILARITY_METRICS_REPR: repr_similarity[1],
+        INTERMEDIATE_SIMILARITY_METRICS_VAL: val_similarity[1]
+    }
+
+    # 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)
+
+    # Save data to a json file.
+    fw_report_utils.dump_report_to_json(report_dir=xquant_config.report_dir,
+                                        collected_data=similarity_metrics)
+
+    return similarity_metrics

mct-nightly-2.1.0.20240618.432/model_compression_toolkit/xquant/common/dataset_utils.py

@@ -0,0 +1,43 @@
+#  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
+
+from model_compression_toolkit.logger import Logger
+
+
+class DatasetUtils:
+    """
+    Class with helpful methods for handling different kinds of datasets from the user.
+    """
+
+    @staticmethod
+    def prepare_dataset(dataset: Callable, 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.
+        """
+
+        Logger.critical("This method should be implemented by the framework-specific DatasetUtils.")  # pragma: no cover
+

mct-nightly-2.1.0.20240618.432/model_compression_toolkit/xquant/common/framework_report_utils.py

@@ -0,0 +1,89 @@
+#  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 json
+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 model_compression_toolkit.xquant.common.constants import REPORT_FILENAME
+from model_compression_toolkit.xquant.common.dataset_utils import DatasetUtils
+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.common.tensorboard_utils import TensorboardUtils
+from model_compression_toolkit.logger import Logger
+
+
+class FrameworkReportUtils:
+    """
+    Class with various utility components required for generating the report in a specific framework.
+    """
+
+    def __init__(self,
+                 fw_info: FrameworkInfo,
+                 fw_impl: FrameworkImplementation,
+                 similarity_calculator: SimilarityCalculator,
+                 dataset_utils: DatasetUtils,
+                 model_folding_utils: ModelFoldingUtils,
+                 tb_utils: TensorboardUtils):
+        """
+        Initializes the FrameworkReportUtils class with various utility components required for generating the report.
+
+        Args:
+            fw_info (FrameworkInfo): Information about the framework being used.
+            fw_impl (FrameworkImplementation): The implemented functions of the framework.
+            similarity_calculator (SimilarityCalculator): A utility for calculating similarity metrics.
+            dataset_utils (DatasetUtils): Utilities for handling datasets.
+            model_folding_utils (ModelFoldingUtils): Utilities for model folding operations.
+            tb_utils (TensorboardUtils): Utilities for TensorBoard operations.
+        """
+        self.fw_info = fw_info
+        self.fw_impl = fw_impl
+        self.similarity_calculator = similarity_calculator
+        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}")
+
+    def dump_report_to_json(self,
+                            report_dir: str,
+                            collected_data: Dict[str, Any]):
+        """
+        Dump the collected data (similarity, etc.) into a JSON file.
+
+        Args:
+            report_dir (str): Directory where the report will be saved.
+            collected_data (Dict[str, Any]): Data collected during report generation.
+
+        """
+        report_file_name = os.path.join(report_dir, REPORT_FILENAME)
+        report_file_name = os.path.abspath(report_file_name)
+        Logger.info(f"Dumping report data to: {report_file_name}")
+
+        with open(report_file_name, 'w') as f:
+            json.dump(collected_data, f, indent=4)

mct-nightly-2.1.0.20240618.432/model_compression_toolkit/xquant/common/model_analyzer.py

@@ -0,0 +1,99 @@
+#  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, Dict, List, Tuple
+
+from abc import ABC, abstractmethod
+
+from model_compression_toolkit.logger import Logger
+
+
+class ModelAnalyzer(ABC):
+    """
+    This class provides abstract methods for analyzing a model, specifically for
+    extracting activations and comparing float and quantized models.
+    """
+
+    @abstractmethod
+    def extract_model_activations(self,
+                                  float_model: Any,
+                                  quantized_model: Any,
+                                  float_name2quant_name: Dict[str, str],
+                                  data: List[Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        """
+        Extracts activations from both the float and quantized models.
+
+        Args:
+            float_model: The float model.
+            quantized_model: The quantized model.
+            float_name2quant_name: A mapping from float model layer names to quantized model layer
+            names.
+            data: Input data for which to compute activations.
+
+        Returns:
+                - Dictionary of activations for the float model.
+                - Dictionary of activations for the quantized model.
+        """
+        Logger.critical("This method should be implemented by the framework-specific ModelAnalyzer.")  # pragma: no cover
+
+
+    @abstractmethod
+    def identify_quantized_compare_points(self, quantized_model: Any) -> 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: The quantized model from which to identify comparison points.
+
+        Returns:
+            List[str]: Names of the layers wrapped with the quantization wrapper.
+        """
+        Logger.critical("This method should be implemented by the framework-specific ModelAnalyzer.")  # pragma: no cover
+
+
+    @abstractmethod
+    def find_corresponding_float_layer(self,
+                                       quant_compare_point: List[str],
+                                       quantized_model: Any) -> str:
+        """
+        Finds the corresponding float model layer for a given quantized model layer.
+
+        Args:
+            quant_compare_point: The name of the quantized model layer.
+            quantized_model: The quantized model.
+
+        Returns:
+            str: The name of the corresponding layer in the float model.
+        """
+        Logger.critical("This method should be implemented by the framework-specific ModelAnalyzer.")  # pragma: no cover
+
+    @abstractmethod
+    def extract_float_layer_names(self, float_model: Any) -> List[str]:
+        """
+        Extracts the names of all layers in the float model.
+
+        Args:
+            float_model: The float model from which to extract layer names.
+
+        Returns:
+            List[str]: Names of all layers in the float model.
+        """
+        Logger.critical("This method should be implemented by the framework-specific ModelAnalyzer.")  # pragma: no cover
+
+

mct-nightly-2.1.0.20240618.432/model_compression_toolkit/xquant/common/model_folding_utils.py

@@ -0,0 +1,104 @@
+#  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.framework_implementation import FrameworkImplementation
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+
+from model_compression_toolkit.core.common.model_builder_mode import ModelBuilderMode
+from model_compression_toolkit.core.common.quantization.quantization_config import DEFAULTCONFIG
+
+from model_compression_toolkit.core.graph_prep_runner import graph_preparation_runner
+from typing import Any, Callable
+
+from model_compression_toolkit.core.common import Graph
+from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
+
+
+class ModelFoldingUtils:
+    """
+    Utility class for handling model folding operations such as batch normalization (BN) folding,
+    residual collapsing, and other graph optimizations.
+    """
+
+    def __init__(self,
+                 fw_info: FrameworkInfo,
+                 fw_impl: FrameworkImplementation,
+                 fw_default_tpc: TargetPlatformCapabilities):
+        """
+        Initialize the ModelFoldingUtils class with framework-specific information, implementation details,
+        and default TPC.
+
+        Args:
+            fw_info: Framework-specific information.
+            fw_impl: Implementation functions for the framework.
+            fw_default_tpc: Default target platform capabilities for the handled framework.
+        """
+        self.fw_info = fw_info
+        self.fw_impl = fw_impl
+        self.fw_default_tpc = fw_default_tpc
+
+    def create_float_folded_model(self, float_model: Any, representative_dataset: Any = None) -> Any:
+        """
+        Create folded version of the model like MCT does (bn folding, residual collapsing, etc.).
+        This is needed since we need the models we compare to have the same architecture for
+        comparing tensors in different points of the models.
+
+        Args:
+            float_model: The floating-point model to be folded.
+            representative_dataset: A callable for generating representative data.
+
+        Returns:
+            The folded floating-point model.
+
+        """
+        float_graph = self.create_float_folded_graph(model=float_model,
+                                                     repr_dataset=representative_dataset)
+        float_folded_model, _ = self.fw_impl.model_builder(
+            float_graph,
+            mode=ModelBuilderMode.FLOAT,
+            append2output=None,
+            fw_info=self.fw_info
+        )
+        return float_folded_model
+
+    def create_float_folded_graph(self, model: Any, repr_dataset: Callable) -> Graph:
+        """
+        Create a folded graph for the float model. This process involves
+        graph optimizations similar to those applied during quantization (e.g., batch normalization folding,
+        residual collapsing).
+
+        Args:
+            model: The floating-point model to be folded into a graph.
+            repr_dataset: A callable that generates representative data.
+
+        Returns:
+            The folded graph.
+        """
+        # TODO:
+        # Consider simplifying graph_preparation_runner by extracting relevant parts to a separate method in MCT.
+        #
+        # Issues:
+        # 1. The quantization config affects how the optimized graph looks (e.g., collapsing).
+        # 2. The back2fw function requires quantization info even for float models.
+        #
+        # Future Considerations:
+        # - Remove quantization config parts related to graph optimizations.
+        # - Update back2fw to handle float models without needing quantization info.
+        graph = graph_preparation_runner(in_model=model,
+                                         representative_data_gen=repr_dataset,
+                                         fw_impl=self.fw_impl,
+                                         fw_info=self.fw_info,
+                                         quantization_config=DEFAULTCONFIG,
+                                         tpc=self.fw_default_tpc)
+        return graph