Remove experimental from pytorch kpi data (#956)

This PR removes the old pytorch_kpi_data method and uses the experimental method as API (after removing 'experimental').

---------

Co-authored-by: reuvenp <reuvenp@altair-semi.com>

model_compression_toolkit/core/__init__.py

@@ -24,4 +24,4 @@
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig, MixedPrecisionQuantizationConfigV2
 from model_compression_toolkit.core.keras.kpi_data_facade import keras_kpi_data, keras_kpi_data_experimental
-from model_compression_toolkit.core.pytorch.kpi_data_facade import pytorch_kpi_data, pytorch_kpi_data_experimental
+from model_compression_toolkit.core.pytorch.kpi_data_facade import pytorch_kpi_data

model_compression_toolkit/core/pytorch/kpi_data_facade.py

@@ -22,8 +22,7 @@
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi_data import compute_kpi_data
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
-from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
-    MixedPrecisionQuantizationConfig, DEFAULT_MIXEDPRECISION_CONFIG, MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfigV2
 from model_compression_toolkit.constants import FOUND_TORCH
 
 if FOUND_TORCH:
@@ -39,73 +38,13 @@
 
     def pytorch_kpi_data(in_model: Module,
                          representative_data_gen: Callable,
-                         quant_config: MixedPrecisionQuantizationConfig = DEFAULT_MIXEDPRECISION_CONFIG,
+                         core_config: CoreConfig = CoreConfig(), # TODO: Why pytorch is initilized and keras not?
                          fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO,
                          target_platform_capabilities: TargetPlatformCapabilities = PYTORCH_DEFAULT_TPC) -> KPI:
         """
         Computes KPI data that can be used to calculate the desired target KPI for mixed-precision quantization.
         Builds the computation graph from the given model and target platform capabilities, and uses it to compute the KPI data.
 
-        Args:
-            in_model (Model): PyTorch model to quantize.
-            representative_data_gen (Callable): Dataset used for calibration.
-            quant_config (MixedPrecisionQuantizationConfig): MixedPrecisionQuantizationConfig containing parameters of how the model should be quantized.
-            fw_info (FrameworkInfo): Information needed for quantization about the specific framework (e.g., kernel channels indices, groups of layers by how they should be quantized, etc.). `Default PyTorch info <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/core/pytorch/default_framework_info.py>`_
-            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
-
-        Returns:
-            A KPI object with total weights parameters sum, max activation tensor and total kpi.
-
-        Examples:
-
-            Import a Pytorch model:
-
-            >>> from torchvision import models
-            >>> module = models.mobilenet_v2()
-
-            Create a random dataset generator, for required number of calibration iterations (num_calibration_batches):
-            In this example a random dataset of 10 batches each containing 4 images is used.
-
-            >>> import numpy as np
-            >>> num_calibration_batches = 10
-            >>> def repr_datagen():
-            >>>     for _ in range(num_calibration_batches):
-            >>>         yield [np.random.random((4, 3, 224, 224))]
-
-            Import mct and call for KPI data calculation:
-
-            >>> import model_compression_toolkit as mct
-            >>> kpi_data = mct.core.pytorch_kpi_data(module, repr_datagen)
-
-        """
-
-        if not isinstance(quant_config, MixedPrecisionQuantizationConfig):
-            Logger.error("KPI data computation can't be executed without MixedPrecisionQuantizationConfig object."
-                         "Given quant_config is not of type MixedPrecisionQuantizationConfig.")
-
-        fw_impl = PytorchImplementation()
-
-        quantization_config, mp_config = quant_config.separate_configs()
-        core_config = CoreConfig(quantization_config=quantization_config,
-                                 mixed_precision_config=mp_config)
-
-        return compute_kpi_data(in_model,
-                                representative_data_gen,
-                                core_config,
-                                target_platform_capabilities,
-                                fw_info,
-                                fw_impl)
-
-
-    def pytorch_kpi_data_experimental(in_model: Module,
-                                      representative_data_gen: Callable,
-                                      core_config: CoreConfig = CoreConfig(),
-                                      fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO,
-                                      target_platform_capabilities: TargetPlatformCapabilities = PYTORCH_DEFAULT_TPC) -> KPI:
-        """
-        Computes KPI data that can be used to calculate the desired target KPI for mixed-precision quantization.
-        Builds the computation graph from the given model and target platform capabilities, and uses it to compute the KPI data.
-
         Args:
             in_model (Model): PyTorch model to quantize.
             representative_data_gen (Callable): Dataset used for calibration.
@@ -155,8 +94,3 @@ def pytorch_kpi_data_experimental(in_model: Module,
     def pytorch_kpi_data(*args, **kwargs):
         Logger.critical('Installing torch is mandatory when using pytorch_kpi_data. '
                         'Could not find Tensorflow package.')  # pragma: no cover
-
-
-    def pytorch_kpi_data_experimental(*args, **kwargs):
-        Logger.critical('Installing torch is mandatory when using pytorch_kpi_data. '
-                        'Could not find Tensorflow package.')  # pragma: no cover

tests/pytorch_tests/function_tests/kpi_data_test.py

@@ -113,8 +113,9 @@ def prep_test(model, mp_bitwidth_candidates_list, random_datagen):
         tpc_name='kpi_data_test')
 
     kpi_data = mct.core.pytorch_kpi_data(in_model=model,
-                                    representative_data_gen=random_datagen,
-                                    target_platform_capabilities=tpc_dict['kpi_data_test'])
+                                         representative_data_gen=random_datagen,
+                                         core_config=mct.core.CoreConfig(mixed_precision_config=mct.core.MixedPrecisionQuantizationConfigV2()),
+                                         target_platform_capabilities=tpc_dict['kpi_data_test'])
 
     return kpi_data
 

tutorials/notebooks/pytorch/ptq/example_pytorch_mobilenet_mixed_precision.py

@@ -111,10 +111,10 @@ def representative_data_gen() -> list:
     # Get KPI information to constraint your model's memory size.
     # Retrieve a KPI object with helpful information of each KPI metric,
     # to constraint the quantized model to the desired memory size.
-    kpi_data = mct.core.pytorch_kpi_data_experimental(model,
-                                                      representative_data_gen,
-                                                      configuration,
-                                                      target_platform_capabilities=target_platform_cap)
+    kpi_data = mct.core.pytorch_kpi_data(model,
+                                         representative_data_gen,
+                                         configuration,
+                                         target_platform_capabilities=target_platform_cap)
 
     # Set a constraint for each of the KPI metrics.
     # Create a KPI object to limit our returned model's size. Note that this values affects only layers and attributes

tutorials/notebooks/pytorch/ptq/example_pytorch_mobilenet_mixed_precision_lut.py

@@ -115,10 +115,10 @@ def representative_data_gen() -> list:
     # Get KPI information to constraint your model's memory size.
     # Retrieve a KPI object with helpful information of each KPI metric,
     # to constraint the quantized model to the desired memory size.
-    kpi_data = mct.core.pytorch_kpi_data_experimental(model,
-                                                      representative_data_gen,
-                                                      configuration,
-                                                      target_platform_capabilities=target_platform_cap)
+    kpi_data = mct.core.pytorch_kpi_data(model,
+                                         representative_data_gen,
+                                         configuration,
+                                         target_platform_capabilities=target_platform_cap)
 
     # Set a constraint for each of the KPI metrics.
     # Create a KPI object to limit our returned model's size. Note that this values affects only layers and attributes

tutorials/notebooks/pytorch/ptq/example_pytorch_mobilenetv2_cifar100_mixed_precision.py

@@ -229,10 +229,10 @@ def representative_data_gen() -> list:
     # Get KPI information to constraint your model's memory size.
     # Retrieve a KPI object with helpful information of each KPI metric,
     # to constraint the quantized model to the desired memory size.
-    kpi_data = mct.core.pytorch_kpi_data_experimental(model,
-                                                      representative_data_gen,
-                                                      configuration,
-                                                      target_platform_capabilities=target_platform_cap)
+    kpi_data = mct.core.pytorch_kpi_data(model,
+                                         representative_data_gen,
+                                         configuration,
+                                         target_platform_capabilities=target_platform_cap)
 
     # Set a constraint for each of the KPI metrics.
     # Create a KPI object to limit our returned model's size. Note that this values affects only layers and attributes

tutorials/quick_start/pytorch_fw/quant.py

@@ -61,7 +61,7 @@ def get_target_kpi(model, weights_compression, representative_data_gen, core_con
         A KPI object computed from MCT and contains info about the target model size.
 
     """
-    kpi_data = mct.core.pytorch_kpi_data_experimental(model, representative_data_gen, core_config=core_config, target_platform_capabilities=tpc)
+    kpi_data = mct.core.pytorch_kpi_data(model, representative_data_gen, core_config=core_config, target_platform_capabilities=tpc)
     weights_kpi = BYTES_TO_FP32 * kpi_data.weights_memory / weights_compression # (4 bytes for fp32) * weights memory(in Bytes) / compression rate
     return KPI(weights_memory=weights_kpi)