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

model_compression_toolkit/xquant/pytorch/dataset_utils.py

@@ -0,0 +1,76 @@
+#  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, Callable
+
+from model_compression_toolkit.xquant.common.dataset_utils import DatasetUtils
+import numpy as np
+
+import torch
+
+
+class PytorchDatasetUtils(DatasetUtils):
+    """
+    Class with helpful methods for handling different kinds of Pytorch 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.
+
+        """
+
+        def process_data(data: Any, is_validation: bool, device: str):
+            """
+            Processes individual data samples: Transfer them to the device, convert to torch tensors if needed,
+            remove labels if this is a validation dataset.
+
+            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.
+            """
+
+            def transfer_to_device(_data):
+                if isinstance(_data, np.ndarray):
+                    return torch.from_numpy(_data).to(device)
+                return _data.to(device)
+
+            if is_validation:
+                inputs = data[0]  # Assume data[0] contains the inputs and data[1] the labels
+                if isinstance(inputs, list):
+                    data = [transfer_to_device(t) for t in inputs]
+                else:
+                    data = [transfer_to_device(inputs)]
+            else:
+                data = [transfer_to_device(t) for t in data]
+
+            return data
+
+        for x in dataset():
+            yield process_data(x, is_validation, device)

model_compression_toolkit/xquant/pytorch/facade_xquant_report.py

@@ -0,0 +1,62 @@
+#  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.constants import FOUND_TORCH
+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_TORCH:
+    from model_compression_toolkit.xquant.pytorch.pytorch_report_utils import PytorchReportUtils
+    import torch
+
+    def xquant_report_pytorch_experimental(float_model: torch.nn.Module,
+                                           quantized_model: torch.nn.Module,
+                                           repr_dataset: Callable,
+                                           validation_dataset: Callable,
+                                           xquant_config: XQuantConfig):
+        """
+        Generate an explainable quantization report for a quantized Pytorch model.
+
+        Args:
+            float_model (torch.nn.Module): The original floating-point Pytorch model.
+            quantized_model (torch.nn.Module): The quantized Pytorch model.
+            repr_dataset (Callable): The representative dataset used during quantization.
+            validation_dataset (Callable): The validation dataset used for evaluation.
+            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)
+
+        pytorch_report_utils = PytorchReportUtils(xquant_config.report_dir)
+
+        _collected_data = core_report_generator(float_model=float_model,
+                                                quantized_model=quantized_model,
+                                                repr_dataset=repr_dataset,
+                                                validation_dataset=validation_dataset,
+                                                fw_report_utils=pytorch_report_utils,
+                                                xquant_config=xquant_config)
+
+        return _collected_data
+
+else:
+    def xquant_report_pytorch_experimental(*args, **kwargs):
+        Logger.critical("PyTorch must be installed to use 'xquant_report_pytorch_experimental'. "
+                                     "The 'torch' package is missing.")  # pragma: no cover

model_compression_toolkit/xquant/pytorch/model_analyzer.py

@@ -0,0 +1,132 @@
+#  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, List, Tuple
+
+import torch
+from mct_quantizers.pytorch.quantize_wrapper import PytorchQuantizationWrapper
+from model_compression_toolkit.xquant.common.constants import MODEL_OUTPUT_KEY
+
+from model_compression_toolkit.xquant.common.model_analyzer import ModelAnalyzer
+
+
+class PytorchModelAnalyzer(ModelAnalyzer):
+    """
+    This class provides utilities for analyzing Pytorch models, specifically for
+    extracting activations and comparing float and quantized models.
+    """
+
+    def extract_model_activations(self,
+                                  float_model: torch.nn.Module,
+                                  quantized_model: torch.nn.Module,
+                                  float_name2quant_name: Dict[str, str],
+                                  data: List[torch.Tensor]) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        """
+        Extracts activations from both the float and quantized models.
+
+        Args:
+            float_model (torch.nn.Module): The float model.
+            quantized_model (torch.nn.Module): The quantized model.
+            float_name2quant_name (Dict[str, str]): A mapping from float model layer names to quantized model layer
+            names.
+            data (List[torch.Tensor]): Input data for which to compute activations.
+
+        Returns:
+            Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+                - Dictionary of activations for the float model.
+                - Dictionary of activations for the quantized model.
+        """
+
+        def _compute_activations(name: str, activations: dict):
+            """
+            Creates a hook function to capture the activations of a layer.
+
+            Args:
+                name (str): The name of the layer.
+                activations (dict): The dictionary to store the activations.
+
+            Returns:
+                hook (function): The hook function to register with the layer.
+            """
+            def hook(model, input, output):
+                activations[name] = output.detach()
+
+            return hook
+
+        # Initialize dictionaries to store activations for both models
+        activations_float = {}
+        activations_quant = {}
+
+        # Register hooks for all layers in the float model
+        for layer_name in float_name2quant_name.keys():
+            layer = dict([*float_model.named_modules()])[layer_name]
+            layer.register_forward_hook(_compute_activations(layer_name, activations_float))
+
+        # Register hooks for all layers in the quantized model
+        for layer_name in float_name2quant_name.values():
+            layer = dict([*quantized_model.named_modules()])[layer_name]
+            layer.register_forward_hook(_compute_activations(layer_name, activations_quant))
+
+        # Perform a forward pass with the input data and capture activations
+        with torch.no_grad():
+            float_predictions = float_model(*data)
+            quant_predictions = quantized_model(*data)
+
+        activations_float[MODEL_OUTPUT_KEY] = float_predictions
+        activations_quant[MODEL_OUTPUT_KEY] = quant_predictions
+
+        return activations_float, activations_quant
+
+    def identify_quantized_compare_points(self,
+                                          quantized_model: torch.nn.Module) -> List[str]:
+        """
+        Identifies points in the quantized model to compare with the float model.
+
+        Args:
+            quantized_model (torch.nn.Module): The quantized model.
+
+        Returns:
+            List[str]: A list of layer names in the quantized model to compare.
+        """
+        return [n for n, m in quantized_model.named_modules() if isinstance(m, PytorchQuantizationWrapper)]
+
+    def find_corresponding_float_layer(self,
+                                       quant_compare_point: str,
+                                       quantized_model: torch.nn.Module) -> str:
+        """
+        Finds the corresponding float model layer for a given quantized model layer.
+        In pytorch, we assume the name is the same in the float model, thus we return quant_compare_point.
+
+        Args:
+            quant_compare_point (str): The name of the layer in the quantized model.
+            quantized_model (torch.nn.Module): The quantized model.
+
+        Returns:
+            str: The name of the corresponding layer in the float model.
+        """
+        return quant_compare_point
+
+    def extract_float_layer_names(self,
+                                  float_model: torch.nn.Module) -> List[str]:
+        """
+        Extracts the names of all layers in the float model.
+
+        Args:
+            float_model (torch.nn.Module): The float model.
+
+        Returns:
+            List[str]: A list of layer names in the float model.
+        """
+        return [n for n, m in float_model.named_modules()]
+

model_compression_toolkit/xquant/pytorch/pytorch_report_utils.py

@@ -0,0 +1,61 @@
+#  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.pytorch.utils import get_working_device
+
+from model_compression_toolkit.ptq.pytorch.quantization_facade import DEFAULT_PYTORCH_TPC
+from model_compression_toolkit.xquant.common.framework_report_utils import FrameworkReportUtils
+from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
+from model_compression_toolkit.core.pytorch.pytorch_implementation import PytorchImplementation
+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.pytorch.dataset_utils import PytorchDatasetUtils
+from model_compression_toolkit.xquant.pytorch.model_analyzer import PytorchModelAnalyzer
+from model_compression_toolkit.xquant.pytorch.similarity_functions import PytorchSimilarityFunctions
+from model_compression_toolkit.xquant.pytorch.tensorboard_utils import PytorchTensorboardUtils
+
+
+class PytorchReportUtils(FrameworkReportUtils):
+    """
+    Class with various utility components required for generating the report for a Pytorch model.
+    """
+    def __init__(self, report_dir: str):
+        """
+        Args:
+            report_dir: Logging dir path.
+        """
+        fw_info = DEFAULT_PYTORCH_INFO
+        fw_impl = PytorchImplementation()
+
+        dataset_utils = PytorchDatasetUtils()
+        model_folding = ModelFoldingUtils(fw_info=fw_info,
+                                          fw_impl=fw_impl,
+                                          fw_default_tpc=DEFAULT_PYTORCH_TPC)
+
+        similarity_calculator = SimilarityCalculator(dataset_utils=dataset_utils,
+                                                     model_folding=model_folding,
+                                                     similarity_functions=PytorchSimilarityFunctions(),
+                                                     model_analyzer_utils=PytorchModelAnalyzer(),
+                                                     device=get_working_device())
+
+        tb_utils = PytorchTensorboardUtils(report_dir=report_dir,
+                                           fw_impl=fw_impl,
+                                           fw_info=fw_info)
+
+        super().__init__(fw_info=fw_info,
+                         fw_impl=fw_impl,
+                         tb_utils=tb_utils,
+                         dataset_utils=dataset_utils,
+                         similarity_calculator=similarity_calculator,
+                         model_folding_utils=model_folding)

model_compression_toolkit/xquant/pytorch/similarity_functions.py

@@ -0,0 +1,68 @@
+#  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.similarity_functions import SimilarityFunctions
+import torch
+
+class PytorchSimilarityFunctions(SimilarityFunctions):
+
+    @staticmethod
+    def compute_mse(x: torch.Tensor, y: torch.Tensor) -> float:
+        """
+        Computes Mean Squared Error between between two tensors (usually, the float and quantized tensors).
+
+        Args:
+            x: Float model predictions.
+            y: Quantized model predictions.
+
+        Returns:
+            Mean Squared Error as a float.
+        """
+        mse = torch.nn.functional.mse_loss(x, y)
+        return mse.item()
+
+    @staticmethod
+    def compute_cs(x: torch.Tensor, y: torch.Tensor) -> float:
+        """
+        Computes Cosine Similarity between two tensors (usually, the float and quantized tensors).
+
+        Args:
+            x: Float model predictions.
+            y: Quantized model predictions.
+
+        Returns:
+            Cosine Similarity as a float.
+        """
+        cs = torch.nn.functional.cosine_similarity(x.flatten(), y.flatten(), dim=0)
+        return cs.item()
+
+    @staticmethod
+    def compute_sqnr(x: torch.Tensor, y: torch.Tensor) -> float:
+        """
+        Computes Signal-to-Quantization-Noise Ratio between two tensors (usually, the float and quantized tensors).
+
+        Args:
+            x: Float model predictions.
+            y: Quantized model predictions.
+
+        Returns:
+            Signal-to-Quantization-Noise Ratio as a float.
+        """
+        signal_power = torch.mean(x ** 2)
+        noise_power = torch.mean((x - y) ** 2)
+        sqnr = signal_power / noise_power
+        return sqnr.item()
+

model_compression_toolkit/xquant/pytorch/tensorboard_utils.py

@@ -0,0 +1,87 @@
+#  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_info import FrameworkInfo
+from typing import Dict, Any, Callable
+
+import torch
+
+from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
+from model_compression_toolkit.core.pytorch.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.model_folding_utils import ModelFoldingUtils
+from model_compression_toolkit.xquant.common.tensorboard_utils import TensorboardUtils
+
+class PytorchTensorboardUtils(TensorboardUtils):
+    """
+    Utility class for handling PyTorch models with TensorBoard. Inherits from TensorboardUtils.
+    This class provides functionalities to display quantized model graphs on TensorBoard.
+    """
+
+    def __init__(self,
+                 report_dir: str,
+                 fw_info: FrameworkInfo,
+                 fw_impl: FrameworkImplementation):
+        """
+        Initialize the PytorchTensorboardUtils instance.
+
+        Args:
+            report_dir: Directory where the reports are stored.
+            fw_info: Information about the framework being used.
+            fw_impl: Implementation methods for the framework.
+        """
+        super().__init__(report_dir,
+                         fw_info,
+                         fw_impl)
+
+    def get_graph_for_tensorboard_display(self,
+                                          quantized_model: torch.nn.Module,
+                                          similarity_metrics: Dict[str, Any],
+                                          repr_dataset: Callable):
+        """
+        Get the graph to display on TensorBoard. The graph represents the quantized model
+        with the similarity metrics that were measured.
+
+        Args:
+            quantized_model: The quantized model to be displayed on TensorBoard.
+            similarity_metrics: Dictionary containing the collected similarity metrics values.
+            repr_dataset: Callable that generates the representative dataset used during graph building.
+
+        Returns:
+            The updated quantized model graph with similarity metrics embedded.
+        """
+        # Read the model and generate a graph representation
+        quant_graph = model_reader(quantized_model,
+                                   representative_data_gen=repr_dataset,
+                                   to_tensor=self.fw_impl.to_tensor,
+                                   to_numpy=self.fw_impl.to_numpy)
+
+        # Iterate through each node in the graph
+        for node in quant_graph.nodes:
+            # Check and add similarity metrics for each node in the graph
+            if node.name in similarity_metrics[INTERMEDIATE_SIMILARITY_METRICS_REPR].keys():
+                node.framework_attr[XQUANT_REPR] = similarity_metrics[INTERMEDIATE_SIMILARITY_METRICS_REPR][f"{node.name}"]
+            elif node.name.removesuffix("_layer") in similarity_metrics[INTERMEDIATE_SIMILARITY_METRICS_REPR].keys():
+                node.framework_attr[XQUANT_REPR] = similarity_metrics[INTERMEDIATE_SIMILARITY_METRICS_REPR][
+                    node.name.removesuffix("_layer")]
+
+            # Check and add validation similarity metrics for each node in the graph
+            if node.name in similarity_metrics[INTERMEDIATE_SIMILARITY_METRICS_VAL].keys():
+                node.framework_attr[XQUANT_VAL] = similarity_metrics[INTERMEDIATE_SIMILARITY_METRICS_VAL][f"{node.name}"]
+            elif node.name.removesuffix("_layer") in similarity_metrics[INTERMEDIATE_SIMILARITY_METRICS_VAL].keys():
+                node.framework_attr[XQUANT_VAL] = similarity_metrics[INTERMEDIATE_SIMILARITY_METRICS_VAL][
+                    node.name.removesuffix("_layer")]
+
+        return quant_graph