Generate report for troubleshooting the quantized model

.github/workflows/run_keras_tests.yml

@@ -31,5 +31,5 @@ jobs:
         run: |
           python -m unittest discover tests/keras_tests/pruning_tests -v
           python -m unittest discover tests/keras_tests/non_parallel_tests -v
-          for script in tests/keras_tests/exporter_tests tests/keras_tests/feature_networks_tests tests/keras_tests/graph_tests tests/keras_tests/layer_tests; do python -m unittest discover $script -v & pids+=($!); done; for pid in ${pids[@]}; do wait $pid || exit 1; done          
+          for script in tests/xquant_tests/keras_tests tests/keras_tests/exporter_tests tests/keras_tests/feature_networks_tests tests/keras_tests/graph_tests tests/keras_tests/layer_tests; do python -m unittest discover $script -v & pids+=($!); done; for pid in ${pids[@]}; do wait $pid || exit 1; done          
 

.github/workflows/run_pytorch_tests.yml

@@ -25,7 +25,9 @@ jobs:
           pip install -r requirements.txt          
           pip install torch==${{ inputs.torch-version }} torchvision onnx onnxruntime
       - name: Run unittests
-        run: python -m unittest discover tests/pytorch_tests -v
+        run: |
+          python -m unittest discover tests/xquant_tests/pytorch_tests -v
+          python -m unittest discover tests/pytorch_tests -v
 
 
 

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
+

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

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'

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

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)