mct-nightly 2.1.0.20240617.451__py3-none-any.whl → 2.1.0.20240619.429__py3-none-any.whl

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.
+#  ==============================================================================
+#

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()]
+