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/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
+

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)

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
+
+

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

model_compression_toolkit/xquant/common/similarity_calculator.py

@@ -0,0 +1,194 @@
+#  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 functools import partial
+
+from typing import Tuple, Any, Dict, Callable
+
+from model_compression_toolkit.xquant.common.constants import MODEL_OUTPUT_KEY
+from model_compression_toolkit.xquant.common.dataset_utils import DatasetUtils
+from model_compression_toolkit.xquant.common.model_analyzer import ModelAnalyzer
+from model_compression_toolkit.xquant.common.model_folding_utils import ModelFoldingUtils
+from model_compression_toolkit.xquant.common.similarity_functions import SimilarityFunctions
+from model_compression_toolkit.logger import Logger
+
+class SimilarityCalculator:
+    """
+    A class to calculate the similarity between two models (that are often referred as float
+    and quantized models). It utilizes various utility classes for dataset preparation, model folding,
+    similarity computation, and model analysis.
+    """
+
+    def __init__(self,
+                 dataset_utils: DatasetUtils,
+                 model_folding: ModelFoldingUtils,
+                 similarity_functions: SimilarityFunctions,
+                 model_analyzer_utils: ModelAnalyzer,
+                 device: str = None):
+        """
+        Initialize the SimilarityCalculator with required utilities.
+
+        Args:
+            dataset_utils (DatasetUtils): Utility class for dataset preparation.
+            model_folding (ModelFoldingUtils): Utility class for model folding operations.
+            similarity_functions (SimilarityFunctions): Class containing similarity functions.
+            model_analyzer_utils (ModelAnalyzer): Utility class for model analysis.
+            device (str, optional): Device to perform computations on (e.g., 'cpu', 'cuda'). Defaults to None.
+        """
+        self.dataset_utils = dataset_utils
+        self.model_folding = model_folding
+        self.similarity_functions = similarity_functions
+        self.model_analyzer_utils = model_analyzer_utils
+        self.device = device
+
+    @staticmethod
+    def compute_tensors_similarity(tensors_to_compare: Tuple[Any, Any],
+                                   similarity_metrics: Dict[str, Callable]) -> Dict[str, float]:
+        """
+        Compute the similarity between two tensors using provided similarity metrics.
+
+        Args:
+            tensors_to_compare (Tuple[Any, Any]): Tensors to compare by computing their similarity.
+            similarity_metrics (Dict[str, Callable]): A dictionary with similarity metric names and functions.
+
+        Returns:
+            Dict[str, float]: A dictionary of similarity metric names and their computed values.
+        """
+        x, y = tensors_to_compare
+        similarity_metrics = {k: v(x, y) for k, v in similarity_metrics.items()}
+        return similarity_metrics
+
+    def _get_float_to_quantized_compare_points(self,
+                                               quantized_model: Any,
+                                               float_model: Any) -> Dict[str, str]:
+        """
+        Map corresponding layers between the float and quantized models for comparison.
+
+        Args:
+            quantized_model (Any): The quantized model.
+            float_model (Any): The float model.
+
+        Returns:
+            Dict[str, str]: A dictionary mapping float model layer names to quantized model layer names.
+        """
+        # Identify the points in the quantized model to compare.
+        quant_points_names = self.model_analyzer_utils.identify_quantized_compare_points(quantized_model)
+
+        float_name2quant_name = {}
+
+        # Extract the names of the layers in the float model.
+        float_layers_names = self.model_analyzer_utils.extract_float_layer_names(float_model)
+
+        # Map each quantized layer to the corresponding float layer.
+        for quant_point in quant_points_names:
+            candidate_float_layer_name = self.model_analyzer_utils.find_corresponding_float_layer(
+                quant_compare_point=quant_point, quantized_model=quantized_model)
+
+            if candidate_float_layer_name in float_layers_names:
+                if candidate_float_layer_name not in float_name2quant_name:
+                    float_name2quant_name[candidate_float_layer_name] = quant_point
+                else:
+                    Logger.critical(f"Duplicate mapping found for layer: {candidate_float_layer_name}.")
+            else:
+                Logger.warning(
+                    f"Could not find a matching layer in the float model for layer with name {quant_point}, "
+                    f"skipping it in similarity metrics comparison points computation.")
+
+        return float_name2quant_name
+
+    def compute_similarity_metrics(self,
+                                   float_model: Any,
+                                   quantized_model: Any,
+                                   dataset: Callable,
+                                   custom_similarity_metrics: Dict[str, Callable] = None,
+                                   is_validation: bool = False) -> Tuple[Dict[str, float], Dict[str, Dict[str, float]]]:
+        """
+        Compute the similarity metrics between the two models (usually, float and quantized models).
+
+        Args:
+            float_model (Any): The float model.
+            quantized_model (Any): The quantized model.
+            dataset (Callable): A callable to provide the dataset.
+            custom_similarity_metrics (Dict[str, Callable], optional): Custom similarity metrics. Defaults to None.
+            is_validation (bool, optional): Flag to indicate if the dataset is for validation. Defaults to False.
+
+        Returns:
+            Tuple[Dict[str, float], Dict[str, Dict[str, float]]]: Aggregated output similarity metrics and
+            intermediate similarity metrics for each layer.
+        """
+        # Prepare the dataset such that the rest of operations are indistinguishable between the representative
+        # dataset and the validation dataset.
+        dataset = partial(self.dataset_utils.prepare_dataset,
+                          dataset=dataset,
+                          is_validation=is_validation,
+                          device=self.device)
+
+        # Create a folded version of the float model.
+        float_model = self.model_folding.create_float_folded_model(float_model=float_model,
+                                                                   representative_dataset=dataset)
+
+        # Gather similarity metrics to compute (default and custom).
+        similarity_metrics_to_compute = self.similarity_functions.get_default_similarity_metrics()
+        if custom_similarity_metrics:
+            if not isinstance(custom_similarity_metrics, dict):
+                Logger.critical(
+                    f"custom_similarity_metrics should be a dictionary but is of type "
+                    f"{type(custom_similarity_metrics)}.")
+            similarity_metrics_to_compute.update(custom_similarity_metrics)
+
+        # Map float model layers to quantized model layers for comparison.
+        float_name2quant_name = self._get_float_to_quantized_compare_points(float_model=float_model,
+                                                                            quantized_model=quantized_model)
+
+        # Initialize dictionaries to store similarity metrics.
+        output_similarity_metrics = {key: [] for key in similarity_metrics_to_compute.keys()}
+        intermediate_similarity_metrics = {layer: {key: [] for key in similarity_metrics_to_compute.keys()} for layer in
+                                           float_name2quant_name.values()}
+
+        # Iterate over the dataset and compute similarity metrics.
+        for x in dataset():
+            # Extract activations and predictions from both models.
+            float_activations, quant_activations = (
+                self.model_analyzer_utils.extract_model_activations(
+                float_model, quantized_model, float_name2quant_name, x))
+
+            float_predictions = float_activations[MODEL_OUTPUT_KEY]
+            quant_predictions = quant_activations[MODEL_OUTPUT_KEY]
+
+            # Compute similarity metrics for the output predictions.
+            output_results = self.compute_tensors_similarity((float_predictions, quant_predictions),
+                                                             similarity_metrics_to_compute)
+            for key in output_similarity_metrics:
+                output_similarity_metrics[key].append(output_results[key])
+
+            # Compute similarity metrics for each intermediate layer.
+            for float_layer, quant_layer in float_name2quant_name.items():
+                intermediate_results = self.compute_tensors_similarity(
+                    (float_activations[float_layer], quant_activations[quant_layer]),
+                    similarity_metrics_to_compute)
+                for key in intermediate_similarity_metrics[quant_layer]:
+                    intermediate_similarity_metrics[quant_layer][key].append(intermediate_results[key])
+
+        # Aggregate the output similarity metrics.
+        aggregated_output_similarity_metrics = {key: sum(value) / len(value) for key, value in
+                                                output_similarity_metrics.items()}
+
+        # Aggregate the intermediate similarity metrics for each layer.
+        for layer_name, layer_similarity_metrics in intermediate_similarity_metrics.items():
+            for similarity_name, similarity_values_list in layer_similarity_metrics.items():
+                if len(similarity_values_list) == 0:
+                    Logger.critical(f"Can not average similarities of an empty list.")
+                intermediate_similarity_metrics[layer_name][similarity_name] = sum(similarity_values_list) / len(similarity_values_list)
+
+        return aggregated_output_similarity_metrics, intermediate_similarity_metrics

model_compression_toolkit/xquant/common/similarity_functions.py

@@ -0,0 +1,81 @@
+#  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, Callable
+
+from model_compression_toolkit.xquant.common.constants import CS_SIMILARITY_METRIC_NAME, SQNR_SIMILARITY_METRIC_NAME, MSE_SIMILARITY_METRIC_NAME
+
+DEFAULT_SIMILARITY_METRICS_NAMES = [CS_SIMILARITY_METRIC_NAME, MSE_SIMILARITY_METRIC_NAME, SQNR_SIMILARITY_METRIC_NAME]
+
+class SimilarityFunctions:
+    """
+    A class that provides various static methods to compute similarity metrics between tensors.
+    """
+
+    @staticmethod
+    def compute_mse(x: Any, y: Any) -> float:
+        """
+        Compute the Mean Squared Error (MSE) between two tensors (usually, the float and quantized predictions).
+
+        Args:
+            x (Any): First tensor to compare.
+            y (Any): Second tensor to compare.
+
+        Returns:
+            float: The computed MSE value.
+        """
+        raise NotImplemented  # pragma: no cover
+
+    @staticmethod
+    def compute_cs(x: Any, y: Any) -> float:
+        """
+        Compute the Cosine Similarity (CS) between two tensors (usually, the float and quantized predictions).
+
+        Args:
+            x (Any): First tensor to compare.
+            y (Any): Second tensor to compare.
+
+        Returns:
+            float: The computed CS value.
+        """
+        raise NotImplemented  # pragma: no cover
+
+    @staticmethod
+    def compute_sqnr(x: Any, y: Any) -> float:
+        """
+        Compute the Signal-to-Quantization-Noise Ratio (SQNR) between two tensors (usually, the float and quantized predictions).
+
+        Args:
+            x (Any): First tensor to compare.
+            y (Any): Second tensor to compare.
+
+        Returns:
+            float: The computed SQNR value.
+        """
+        raise NotImplemented  # pragma: no cover
+
+    def get_default_similarity_metrics(self) -> Dict[str, Callable]:
+        """
+        Get the default similarity metrics to compute.
+
+        Returns:
+            Dict[str, Callable]: A dictionary where the keys are similarity metric names and the values are the corresponding functions.
+        """
+        return {
+            MSE_SIMILARITY_METRIC_NAME: self.compute_mse,
+            CS_SIMILARITY_METRIC_NAME: self.compute_cs,
+            SQNR_SIMILARITY_METRIC_NAME: self.compute_sqnr
+        }
+