mct-nightly 1.11.0.20240317.91316__tar.gz → 1.11.0.20240319.407__tar.gz

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 1.11.0.20240317.91316
+Version: 1.11.0.20240319.407
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/mct_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 1.11.0.20240317.91316
+Version: 1.11.0.20240319.407
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/__init__.py

@@ -27,4 +27,4 @@ from model_compression_toolkit import data_generation
 from model_compression_toolkit import pruning
 from model_compression_toolkit.trainable_infrastructure.keras.load_model import keras_load_quantized_model
 
-__version__ = "1.11.0.20240317.091316"
+__version__ = "1.11.0.20240319.000407"

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/core/__init__.py

@@ -25,3 +25,5 @@ from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core.keras.kpi_data_facade import keras_kpi_data
 from model_compression_toolkit.core.pytorch.kpi_data_facade import pytorch_kpi_data
+from model_compression_toolkit.core.common.mixed_precision.distance_weighting import MpDistanceWeighting
+

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/core/common/mixed_precision/kpi_tools/kpi_data.py

@@ -47,13 +47,15 @@ def compute_kpi_data(in_model: Any,
 
     """
 
+    # We assume that the kpi_data API is used to compute the model KPI for mixed precision scenario,
+    # so we run graph preparation under the assumption of enabled mixed precision.
     transformed_graph = graph_preparation_runner(in_model,
                                                  representative_data_gen,
                                                  core_config.quantization_config,
                                                  fw_info,
                                                  fw_impl,
                                                  tpc,
-                                                 mixed_precision_enable=core_config.mixed_precision_enable)
+                                                 mixed_precision_enable=True)
 
     # Compute parameters sum
     weights_params = compute_nodes_weights_params(graph=transformed_graph, fw_info=fw_info)

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py

@@ -16,13 +16,11 @@
 from typing import List, Callable
 
 from model_compression_toolkit.core.common.mixed_precision.distance_weighting import MpDistanceWeighting
-from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
 
 
 class MixedPrecisionQuantizationConfig:
 
     def __init__(self,
-                 target_kpi: KPI = None,
                  compute_distance_fn: Callable = None,
                  distance_weighting_method: MpDistanceWeighting = MpDistanceWeighting.AVG,
                  num_of_images: int = 32,
@@ -36,7 +34,6 @@ class MixedPrecisionQuantizationConfig:
         Class with mixed precision parameters to quantize the input model.
 
         Args:
-            target_kpi (KPI): KPI to constraint the search of the mixed-precision configuration for the model.
             compute_distance_fn (Callable): Function to compute a distance between two tensors. If None, using pre-defined distance methods based on the layer type for each layer.
             distance_weighting_method (MpDistanceWeighting): MpDistanceWeighting enum value that provides a function to use when weighting the distances among different layers when computing the sensitivity metric.
             num_of_images (int): Number of images to use to evaluate the sensitivity of a mixed-precision model comparing to the float model.
@@ -49,7 +46,6 @@ class MixedPrecisionQuantizationConfig:
 
         """
 
-        self.target_kpi = target_kpi
         self.compute_distance_fn = compute_distance_fn
         self.distance_weighting_method = distance_weighting_method
         self.num_of_images = num_of_images
@@ -67,13 +63,21 @@ class MixedPrecisionQuantizationConfig:
 
         self.metric_normalization_threshold = metric_normalization_threshold
 
-    def set_target_kpi(self, target_kpi: KPI):
+        self._mixed_precision_enable = False
+
+    def set_mixed_precision_enable(self):
+        """
+        Set a flag in mixed precision config indicating that mixed precision is enabled.
         """
-        Setting target KPI in mixed precision config.
 
-        Args:
-            target_kpi: A target KPI to set.
+        self._mixed_precision_enable = True
 
+    @property
+    def mixed_precision_enable(self):
         """
+        A property that indicates whether mixed precision quantization is enabled.
 
-        self.target_kpi = target_kpi
+        Returns: True if mixed precision quantization is enabled
+
+        """
+        return self._mixed_precision_enable

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py

@@ -47,6 +47,7 @@ search_methods = {
 def search_bit_width(graph_to_search_cfg: Graph,
                      fw_info: FrameworkInfo,
                      fw_impl: FrameworkImplementation,
+                     target_kpi: KPI,
                      mp_config: MixedPrecisionQuantizationConfig,
                      representative_data_gen: Callable,
                      search_method: BitWidthSearchMethod = BitWidthSearchMethod.INTEGER_PROGRAMMING,
@@ -63,6 +64,7 @@ def search_bit_width(graph_to_search_cfg: Graph,
         graph_to_search_cfg: Graph to search a MP configuration for.
         fw_info: FrameworkInfo object about the specific framework (e.g., attributes of different layers' weights to quantize).
         fw_impl: FrameworkImplementation object with specific framework methods implementation.
+        target_kpi: Target KPI to bound our feasible solution space s.t the configuration does not violate it.
         mp_config: Mixed-precision quantization configuration.
         representative_data_gen: Dataset to use for retrieving images for the models inputs.
         search_method: BitWidthSearchMethod to define which searching method to use.
@@ -74,7 +76,6 @@ def search_bit_width(graph_to_search_cfg: Graph,
         bit-width index on the node).
 
     """
-    target_kpi = mp_config.target_kpi
 
     # target_kpi have to be passed. If it was not passed, the facade is not supposed to get here by now.
     if target_kpi is None:

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/core/common/pruning/pruning_config.py

@@ -20,14 +20,20 @@ from model_compression_toolkit.constants import PRUNING_NUM_SCORE_APPROXIMATIONS
 
 class ImportanceMetric(Enum):
     """
-    Enum for specifying the metric used to determine the importance of channels when pruning.
+    Enum for specifying the metric used to determine the importance of channels when pruning:
+
+    LFH - Label-Free Hessian uses hessian info for measuring each channel's sensitivity.
+
     """
     LFH = 0  # Score based on the Hessian matrix w.r.t. layers weights, to determine channel importance without labels.
 
 
 class ChannelsFilteringStrategy(Enum):
     """
-    Enum for specifying the strategy used for filtering (pruning) channels.
+    Enum for specifying the strategy used for filtering (pruning) channels:
+
+    GREEDY - Prune the least important channel groups up to allowed resources in the KPI (for now, only weights_memory is considered).
+
     """
     GREEDY = 0  # Greedy strategy for pruning channels based on importance metrics.
 

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/core/common/pruning/pruning_info.py

@@ -26,23 +26,16 @@ class PruningInfo:
     and importance scores for each layer. This class acts as a container for accessing
     pruning-related metadata.
 
-    Attributes:
-        pruning_masks (Dict[BaseNode, np.ndarray]): Stores the pruning masks for each layer.
-            A pruning mask is an array where each element indicates whether the corresponding
-            channel or neuron has been pruned (0) or kept (1).
-        importance_scores (Dict[BaseNode, np.ndarray]): Stores the importance scores for each layer.
-            Importance scores quantify the significance of each channel in the layer.
     """
 
     def __init__(self,
                  pruning_masks: Dict[BaseNode, np.ndarray],
                  importance_scores: Dict[BaseNode, np.ndarray]):
         """
-        Initializes the PruningInfo with pruning masks and importance scores.
-
         Args:
-            pruning_masks (Dict[BaseNode, np.ndarray]): Pruning masks for each layer.
-            importance_scores (Dict[BaseNode, np.ndarray]): Importance scores for each layer.
+            pruning_masks (Dict[BaseNode, np.ndarray]): Stores the pruning masks for each layer. A pruning mask is an array where each element indicates whether the corresponding channel or neuron has been pruned (0) or kept (1).
+            importance_scores (Dict[BaseNode, np.ndarray]): Stores the importance scores for each layer. Importance scores quantify the significance of each channel in the layer.
+
         """
         self._pruning_masks = pruning_masks
         self._importance_scores = importance_scores

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/core/common/quantization/core_config.py

@@ -30,14 +30,19 @@ class CoreConfig:
 
         Args:
             quantization_config (QuantizationConfig): Config for quantization.
-            mixed_precision_config (MixedPrecisionQuantizationConfig): Config for mixed precision quantization (optional, default=None).
+            mixed_precision_config (MixedPrecisionQuantizationConfig): Config for mixed precision quantization.
+            If None, a default MixedPrecisionQuantizationConfig is used.
             debug_config (DebugConfig): Config for debugging and editing the network quantization process.
         """
         self.quantization_config = quantization_config
-        self.mixed_precision_config = mixed_precision_config
         self.debug_config = debug_config
 
+        if mixed_precision_config is None:
+            self.mixed_precision_config = MixedPrecisionQuantizationConfig()
+        else:
+            self.mixed_precision_config = mixed_precision_config
+
     @property
     def mixed_precision_enable(self):
-        return self.mixed_precision_config is not None
+        return self.mixed_precision_config is not None and self.mixed_precision_config.mixed_precision_enable
 

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/core/common/quantization/set_node_quantization_config.py

@@ -71,7 +71,7 @@ def set_quantization_configs_to_node(node: BaseNode,
         quant_config: Quantization configuration to generate the node's configurations from.
         fw_info: Information needed for quantization about the specific framework.
         tpc: TargetPlatformCapabilities to get default OpQuantizationConfig.
-        mixed_precision_enable: is mixed precision enabled
+        mixed_precision_enable: is mixed precision enabled.
     """
     node_qc_options = node.get_qco(tpc)
 

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/core/graph_prep_runner.py

@@ -57,7 +57,8 @@ def graph_preparation_runner(in_model: Any,
         fw_impl: FrameworkImplementation object with a specific framework methods implementation.
         tpc: TargetPlatformCapabilities object that models the inference target platform and
             the attached framework operator's information.
-        tb_w: TensorboardWriter object for logging
+        tb_w: TensorboardWriter object for logging.
+        mixed_precision_enable: is mixed precision enabled.
 
     Returns:
         An internal graph representation of the input model.
@@ -103,7 +104,7 @@ def get_finalized_graph(initial_graph: Graph,
             kernel channels indices, groups of layers by how they should be quantized, etc.)
         tb_w (TensorboardWriter): TensorboardWriter object to use for logging events such as graphs, histograms, etc.
         fw_impl (FrameworkImplementation): FrameworkImplementation object with a specific framework methods implementation.
-            mixed_precision_enable: is mixed precision enabled.
+        mixed_precision_enable: is mixed precision enabled.
 
     Returns: Graph object that represents the model, after applying all required modifications to it.
     """

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/core/pytorch/kpi_data_facade.py

@@ -38,7 +38,7 @@ if FOUND_TORCH:
 
     def pytorch_kpi_data(in_model: Module,
                          representative_data_gen: Callable,
-                         core_config: CoreConfig = CoreConfig(mixed_precision_config=MixedPrecisionQuantizationConfig()),
+                         core_config: CoreConfig = CoreConfig(),
                          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.

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/core/runner.py

@@ -47,6 +47,7 @@ def core_runner(in_model: Any,
                 fw_info: FrameworkInfo,
                 fw_impl: FrameworkImplementation,
                 tpc: TargetPlatformCapabilities,
+                target_kpi: KPI = None,
                 tb_w: TensorboardWriter = None):
     """
     Quantize a trained model using post-training quantization.
@@ -66,6 +67,7 @@ def core_runner(in_model: Any,
         fw_impl: FrameworkImplementation object with a specific framework methods implementation.
         tpc: TargetPlatformCapabilities object that models the inference target platform and
                                               the attached framework operator's information.
+        target_kpi: KPI to constraint the search of the mixed-precision configuration for the model.
         tb_w: TensorboardWriter object for logging
 
     Returns:
@@ -81,6 +83,13 @@ def core_runner(in_model: Any,
         Logger.warning('representative_data_gen generates a batch size of 1 which can be slow for optimization:'
                        ' consider increasing the batch size')
 
+    # Checking whether to run mixed precision quantization
+    if target_kpi is not None:
+        if core_config.mixed_precision_config is None:
+            Logger.critical("Provided an initialized target_kpi, that means that mixed precision quantization is "
+                            "enabled, but the provided MixedPrecisionQuantizationConfig is None.")
+        core_config.mixed_precision_config.set_mixed_precision_enable()
+
     graph = graph_preparation_runner(in_model,
                                      representative_data_gen,
                                      core_config.quantization_config,
@@ -105,13 +114,12 @@ def core_runner(in_model: Any,
     # Finalize bit widths
     ######################################
     if core_config.mixed_precision_enable:
-        if core_config.mixed_precision_config.target_kpi is None:
-            Logger.critical(f"Trying to run Mixed Precision quantization without providing a valid target KPI.")
         if core_config.mixed_precision_config.configuration_overwrite is None:
 
             bit_widths_config = search_bit_width(tg,
                                                  fw_info,
                                                  fw_impl,
+                                                 target_kpi,
                                                  core_config.mixed_precision_config,
                                                  representative_data_gen,
                                                  hessian_info_service=hessian_info_service)

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/data_generation/__init__.py

@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+
 from model_compression_toolkit.constants import FOUND_TORCH, FOUND_TF
+from model_compression_toolkit.data_generation.common.data_generation_config import DataGenerationConfig
+from model_compression_toolkit.data_generation.common.enums import ImageGranularity, DataInitType, SchedulerType, BNLayerWeightingType, OutputLossType, BatchNormAlignemntLossType, ImagePipelineType, ImageNormalizationType
 
 if FOUND_TF:
     from model_compression_toolkit.data_generation.keras.keras_data_generation import (

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/data_generation/common/enums.py

@@ -29,10 +29,14 @@ class EnumBaseClass(Enum):
 
 class ImageGranularity(EnumBaseClass):
     """
-    An enum for choosing the image dependence granularity when generating images.
-    0. ImageWise
-    1. BatchWise
-    2. AllImages
+    An enum for choosing the image dependence granularity when generating images:
+
+    ImageWise
+
+    BatchWise
+
+    AllImages
+
     """
 
     ImageWise = 0
@@ -42,9 +46,12 @@ class ImageGranularity(EnumBaseClass):
 
 class DataInitType(EnumBaseClass):
     """
-    An enum for choosing the image dependence granularity when generating images.
-    0. Gaussian
-    1. Diverse
+    An enum for choosing the image dependence granularity when generating images:
+
+    Gaussian
+
+    Diverse
+
     """
     Gaussian = 0
     Diverse = 1
@@ -52,9 +59,14 @@ class DataInitType(EnumBaseClass):
 
 class ImagePipelineType(EnumBaseClass):
     """
-    An enum for choosing the image pipeline type for image manipulation.
-    RANDOM_CROP_FLIP: Crop and flip the images.
-    IDENTITY: Do not apply any manipulation (identity transformation).
+    An enum for choosing the image pipeline type for image manipulation:
+
+    RANDOM_CROP - Crop the images.
+
+    RANDOM_CROP_FLIP - Crop and flip the images.
+
+    IDENTITY - Do not apply any manipulation (identity transformation).
+
     """
     RANDOM_CROP = 'random_crop'
     RANDOM_CROP_FLIP = 'random_crop_flip'
@@ -63,10 +75,14 @@ class ImagePipelineType(EnumBaseClass):
 
 class ImageNormalizationType(EnumBaseClass):
     """
-    An enum for choosing the image normalization type.
-    TORCHVISION: Normalize the images using torchvision normalization.
-    KERAS_APPLICATIONS: Normalize the images using keras_applications imagenet normalization.
-    NO_NORMALIZATION: Do not apply any normalization.
+    An enum for choosing the image normalization type:
+
+    TORCHVISION - Normalize the images using torchvision normalization.
+
+    KERAS_APPLICATIONS - Normalize the images using keras_applications imagenet normalization.
+
+    NO_NORMALIZATION - Do not apply any normalization.
+
     """
     TORCHVISION = 'torchvision'
     KERAS_APPLICATIONS = 'keras_applications'
@@ -75,10 +91,14 @@ class ImageNormalizationType(EnumBaseClass):
 
 class BNLayerWeightingType(EnumBaseClass):
     """
-   An enum for choosing the layer weighting type.
-   AVERAGE: Use the same weight per layer.
-   FIRST_LAYER_MULTIPLIER: Use a multiplier for the first layer, all other layers with the same weight.
-   GRAD: Use gradient-based layer weighting.
+   An enum for choosing the layer weighting type:
+
+   AVERAGE - Use the same weight per layer.
+
+   FIRST_LAYER_MULTIPLIER - Use a multiplier for the first layer, all other layers with the same weight.
+
+   GRAD - Use gradient-based layer weighting.
+
    """
     AVERAGE = 'average'
     FIRST_LAYER_MULTIPLIER = 'first_layer_multiplier'
@@ -87,18 +107,24 @@ class BNLayerWeightingType(EnumBaseClass):
 
 class BatchNormAlignemntLossType(EnumBaseClass):
     """
-    An enum for choosing the BatchNorm alignment loss type.
-    L2_SQUARE: Use L2 square loss for BatchNorm alignment.
+    An enum for choosing the BatchNorm alignment loss type:
+
+    L2_SQUARE - Use L2 square loss for BatchNorm alignment.
+
     """
     L2_SQUARE = 'l2_square'
 
 
 class OutputLossType(EnumBaseClass):
     """
-    An enum for choosing the output loss type.
-    NONE: No output loss is applied.
-    MIN_MAX_DIFF: Use min-max difference as the output loss.
-    REGULARIZED_MIN_MAX_DIFF: Use regularized min-max difference as the output loss.
+    An enum for choosing the output loss type:
+
+    NONE - No output loss is applied.
+
+    MIN_MAX_DIFF - Use min-max difference as the output loss.
+
+    REGULARIZED_MIN_MAX_DIFF - Use regularized min-max difference as the output loss.
+
     """
     NONE = 'none'
     MIN_MAX_DIFF = 'min_max_diff'
@@ -107,9 +133,12 @@ class OutputLossType(EnumBaseClass):
 
 class SchedulerType(EnumBaseClass):
     """
-    An enum for choosing the scheduler type for the optimizer.
-    REDUCE_ON_PLATEAU: Use the ReduceOnPlateau scheduler.
-    STEP: Use the Step scheduler.
+    An enum for choosing the scheduler type for the optimizer:
+
+    REDUCE_ON_PLATEAU - Use the ReduceOnPlateau scheduler.
+
+    STEP - Use the Step scheduler.
+
     """
     REDUCE_ON_PLATEAU = 'reduce_on_plateau'
     STEP = 'step'

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/data_generation/keras/keras_data_generation.py

@@ -131,7 +131,36 @@ if FOUND_TF:
 
         Returns:
             List[tf.Tensor]: Finalized list containing generated images.
+
+        Examples:
+
+            In this example, we'll walk through generating images using a simple Keras model and a data generation configuration. The process involves creating a model, setting up a data generation configuration, and finally generating images with specified parameters.
+
+            Start by importing the Model Compression Toolkit (MCT), TensorFlow, and some layers from `tensorflow.keras`:
+
+            >>> import model_compression_toolkit as mct
+            >>> from tensorflow.keras.models import Sequential
+            >>> from tensorflow.keras.layers import Conv2D, BatchNormalization, Flatten, Dense, Reshape
+
+            Next, define a simple Keras model:
+
+            >>> model = Sequential([Conv2D(2, 3, input_shape=(8,8,3)), BatchNormalization(), Flatten(), Dense(10)])
+
+            Configure the data generation process using `get_keras_data_generation_config`. This function allows customization of the data generation process. For simplicity, this example sets the number of iterations (`n_iter`) to 1 and the batch size (`data_gen_batch_size`) to 2.
+
+            >>> config = mct.data_generation.get_keras_data_generation_config(n_iter=1, data_gen_batch_size=2)
+
+            Finally, use the `keras_data_generation_experimental` function to generate images based on the model and data generation configuration.
+            Notice that this function is experimental and may change in future versions of MCT.
+            The `n_images` parameter specifies the number of images to generate, and `output_image_size` sets the size of the generated images.
+
+            >>> generated_images = mct.data_generation.keras_data_generation_experimental(model=model, n_images=4, output_image_size=(8, 8), data_generation_config=config)
+
+            The generated images can then be used for various purposes, such as data-free quantization.
+
+
         """
+
         Logger.warning(f"keras_data_generation_experimental is experimental "
                        f"and is subject to future changes."
                        f"If you encounter an issue, please open an issue in our GitHub "

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/data_generation/pytorch/pytorch_data_generation.py

@@ -129,7 +129,7 @@ if FOUND_TORCH:
     def pytorch_data_generation_experimental(
             model: Module,
             n_images: int,
-            output_image_size: Tuple,
+            output_image_size: int,
             data_generation_config: DataGenerationConfig) -> List[Tensor]:
         """
         Function to perform data generation using the provided model and data generation configuration.
@@ -137,11 +137,38 @@ if FOUND_TORCH:
         Args:
             model (Module): PyTorch model to generate data for.
             n_images (int): Number of images to generate.
-            output_image_size (Tuple): Size of the output images.
+            output_image_size (int): The hight and width size of the output images.
             data_generation_config (DataGenerationConfig): Configuration for data generation.
 
         Returns:
             List[Tensor]: Finalized list containing generated images.
+
+        Examples:
+
+            In this example, we'll walk through generating images using a simple PyTorch model and a data generation configuration. The process involves creating a model, setting up a data generation configuration, and finally generating images with specified parameters.
+
+            Start by importing the Model Compression Toolkit (MCT), PyTorch, and some modules from `torch.nn`:
+
+            >>> import model_compression_toolkit as mct
+            >>> import torch.nn as nn
+            >>> from torch.nn import Conv2d, BatchNorm2d, Flatten, Linear
+
+            Next, define a simple PyTorch model:
+
+            >>> model = nn.Sequential(nn.Conv2d(3, 2, 3), nn.BatchNorm2d(2), nn.Flatten(), nn.Linear(2*6*6, 10))
+
+            Configure the data generation process using `get_pytorch_data_generation_config`. This function allows customization of the data generation process. For simplicity, this example sets the number of iterations (`n_iter`) to 1 and the batch size (`data_gen_batch_size`) to 2.
+
+            >>> config = mct.data_generation.get_pytorch_data_generation_config(n_iter=1, data_gen_batch_size=2)
+
+            Finally, use the `pytorch_data_generation_experimental` function to generate images based on the model and data generation configuration.
+            Notice that this function is experimental and may change in future versions of MCT.
+            The `n_images` parameter specifies the number of images to generate, and `output_image_size` sets the size of the generated images.
+
+            >>> generated_images = mct.data_generation.pytorch_data_generation_experimental(model=model, n_images=4, output_image_size=8, data_generation_config=config)
+
+            The generated images can then be used for various purposes, such as data-free quantization.
+
         """
 
         Logger.warning(f"pytorch_data_generation_experimental is experimental "

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/exporter/model_exporter/fw_agonstic/quantization_format.py

@@ -16,6 +16,16 @@ from enum import Enum
 
 
 class QuantizationFormat(Enum):
+    """
+    Specify which quantization format to use for exporting a quantized model.
+
+    FAKELY_QUANT - Weights and activations are quantized but represented using float data type.
+
+    INT8 - Weights and activations are represented using 8-bit integer data type.
+
+    MCTQ - Weights and activations are quantized using mct_quantizers custom quantizers.
+
+    """
     FAKELY_QUANT = 0
     INT8 = 1
     MCTQ = 2

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/exporter/model_exporter/keras/keras_export_facade.py

@@ -42,21 +42,17 @@ if FOUND_TF:
                            serialization_format: KerasExportSerializationFormat = KerasExportSerializationFormat.KERAS,
                            quantization_format : QuantizationFormat = QuantizationFormat.MCTQ) -> Dict[str, type]:
         """
-        Export a Keras quantized model to a h5 or tflite model.
+        Export a Keras quantized model to a .keras or .tflite format model (according to serialization_format).
         The model will be saved to the path in save_model_path.
-        keras_export_model supports the combination of QuantizationFormat.FAKELY_QUANT (where weights
-        and activations are float fakely-quantized values) and KerasExportSerializationFormat.KERAS_H5 (where the model
-        will be saved to h5 model) or the combination of KerasExportSerializationFormat.TFLITE (where the model will be
-        saved to tflite model) with QuantizationFormat.FAKELY_QUANT or QuantizationFormat.INT8 (where weights and
-        activations are represented using 8bits integers).
+        Models that are exported to .keras format can use quantization_format of QuantizationFormat.MCTQ or QuantizationFormat.FAKELY_QUANT.
+        Models that are exported to .tflite format can use quantization_format of QuantizationFormat.INT8 or QuantizationFormat.FAKELY_QUANT.
 
         Args:
             model: Model to export.
             save_model_path: Path to save the model.
             is_layer_exportable_fn: Callable to check whether a layer can be exported or not.
-            serialization_format: Format to export the model according to (by default
-            KerasExportSerializationFormat.KERAS_H5).
-            quantization_format: Format of how quantizers are exported (fakely-quant, int8, MCTQ quantizers).
+            serialization_format: Format to export the model according to (KerasExportSerializationFormat.KERAS, by default).
+            quantization_format: Format of how quantizers are exported (MCTQ quantizers, by default).
 
         Returns:
             Custom objects dictionary needed to load the model.

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/gptq/common/gptq_config.py

@@ -19,9 +19,12 @@ from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT
 
 class RoundingType(Enum):
     """
-    An enum for choosing the GPTQ rounding methods
-    0. STRAIGHT-THROUGH ESTIMATOR
-    1. SoftQuantizer
+    An enum for choosing the GPTQ rounding methods:
+
+    STE - STRAIGHT-THROUGH ESTIMATOR
+
+    SoftQuantizer - SoftQuantizer
+
     """
     STE = 0
     SoftQuantizer = 1

{mct-nightly-1.11.0.20240317.91316 → mct-nightly-1.11.0.20240319.407}/model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -116,6 +116,7 @@ if FOUND_TF:
     def keras_gradient_post_training_quantization(in_model: Model, representative_data_gen: Callable,
                                                   gptq_config: GradientPTQConfig,
                                                   gptq_representative_data_gen: Callable = None,
+                                                  target_kpi: KPI = None,
                                                   core_config: CoreConfig = CoreConfig(),
                                                   target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC) -> Tuple[Model, UserInformation]:
         """
@@ -139,6 +140,7 @@ if FOUND_TF:
             representative_data_gen (Callable): Dataset used for calibration.
             gptq_config (GradientPTQConfig): Configuration for using gptq (e.g. optimizer).
             gptq_representative_data_gen (Callable): Dataset used for GPTQ training. If None defaults to representative_data_gen
+            target_kpi (KPI): KPI object to limit the search of the mixed-precision configuration as desired.
             core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
             target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
 
@@ -166,6 +168,12 @@ if FOUND_TF:
 
             >>> config = mct.core.CoreConfig()
 
+            If mixed precision is desired, create an MCT core config with a mixed-precision configuration, to quantize a model
+            with different bitwidths for different layers.
+            The candidates bitwidth for quantization should be defined in the target platform model:
+
+            >>> config = mct.core.CoreConfig(mixed_precision_config=mct.core.MixedPrecisionQuantizationConfig(num_of_images=1))
+
             For mixed-precision set a target KPI object:
             Create a KPI object to limit our returned model's size. Note that this value affects only coefficients
             that should be quantized (for example, the kernel of Conv2D in Keras will be affected by this value,
@@ -173,19 +181,13 @@ if FOUND_TF:
 
             >>> kpi = mct.core.KPI(model.count_params() * 0.75)  # About 0.75 of the model size when quantized with 8 bits.
 
-            If mixed precision is desired, create an MCT core config with a mixed-precision configuration, to quantize a model
-            with different bitwidths for different layers.
-            The candidates bitwidth for quantization should be defined in the target platform model:
-
-            >>> config = mct.core.CoreConfig(mixed_precision_config=mct.core.MixedPrecisionQuantizationConfig(num_of_images=1, target_kpi=kpi))
-
             Create GPTQ config:
 
             >>> gptq_config = mct.gptq.get_keras_gptq_config(n_epochs=1)
 
             Pass the model with the representative dataset generator to get a quantized model:
 
-            >>> quantized_model, quantization_info = mct.gptq.keras_gradient_post_training_quantization(model, repr_datagen, gptq_config, core_config=config)
+            >>> quantized_model, quantization_info = mct.gptq.keras_gradient_post_training_quantization(model, repr_datagen, gptq_config, target_kpi=kpi, core_config=config)
 
         """
         KerasModelValidation(model=in_model,