mct-nightly 2.2.0.20241026.508__tar.gz → 2.2.0.20241027.532__tar.gz

{mct-nightly-2.2.0.20241026.508 → mct-nightly-2.2.0.20241027.532}/PKG-INFO

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

{mct-nightly-2.2.0.20241026.508 → mct-nightly-2.2.0.20241027.532}/mct_nightly.egg-info/PKG-INFO

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

{mct-nightly-2.2.0.20241026.508 → mct-nightly-2.2.0.20241027.532}/mct_nightly.egg-info/SOURCES.txt

@@ -360,6 +360,8 @@ model_compression_toolkit/gptq/common/gptq_constants.py
 model_compression_toolkit/gptq/common/gptq_framework_implementation.py
 model_compression_toolkit/gptq/common/gptq_graph.py
 model_compression_toolkit/gptq/common/gptq_training.py
+model_compression_toolkit/gptq/common/gradual_activation_quantization.py
+model_compression_toolkit/gptq/common/regularization_factory.py
 model_compression_toolkit/gptq/keras/__init__.py
 model_compression_toolkit/gptq/keras/gptq_keras_implementation.py
 model_compression_toolkit/gptq/keras/gptq_loss.py
@@ -370,7 +372,6 @@ model_compression_toolkit/gptq/keras/quantizer/__init__.py
 model_compression_toolkit/gptq/keras/quantizer/base_keras_gptq_quantizer.py
 model_compression_toolkit/gptq/keras/quantizer/quant_utils.py
 model_compression_toolkit/gptq/keras/quantizer/quantization_builder.py
-model_compression_toolkit/gptq/keras/quantizer/regularization_factory.py
 model_compression_toolkit/gptq/keras/quantizer/soft_rounding/__init__.py
 model_compression_toolkit/gptq/keras/quantizer/soft_rounding/soft_quantizer_reg.py
 model_compression_toolkit/gptq/keras/quantizer/soft_rounding/symmetric_soft_quantizer.py
@@ -385,10 +386,8 @@ model_compression_toolkit/gptq/pytorch/graph_info.py
 model_compression_toolkit/gptq/pytorch/quantization_facade.py
 model_compression_toolkit/gptq/pytorch/quantizer/__init__.py
 model_compression_toolkit/gptq/pytorch/quantizer/base_pytorch_gptq_quantizer.py
-model_compression_toolkit/gptq/pytorch/quantizer/gradual_activation_quantization.py
 model_compression_toolkit/gptq/pytorch/quantizer/quant_utils.py
 model_compression_toolkit/gptq/pytorch/quantizer/quantization_builder.py
-model_compression_toolkit/gptq/pytorch/quantizer/regularization_factory.py
 model_compression_toolkit/gptq/pytorch/quantizer/soft_rounding/__init__.py
 model_compression_toolkit/gptq/pytorch/quantizer/soft_rounding/soft_quantizer_reg.py
 model_compression_toolkit/gptq/pytorch/quantizer/soft_rounding/symmetric_soft_quantizer.py
@@ -502,6 +501,7 @@ model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/v1/
 model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/v1/tpc_pytorch.py
 model_compression_toolkit/trainable_infrastructure/__init__.py
 model_compression_toolkit/trainable_infrastructure/common/__init__.py
+model_compression_toolkit/trainable_infrastructure/common/annealing_schedulers.py
 model_compression_toolkit/trainable_infrastructure/common/base_trainable_quantizer.py
 model_compression_toolkit/trainable_infrastructure/common/constants.py
 model_compression_toolkit/trainable_infrastructure/common/get_quantizer_config.py
@@ -509,7 +509,9 @@ model_compression_toolkit/trainable_infrastructure/common/get_quantizers.py
 model_compression_toolkit/trainable_infrastructure/common/quant_utils.py
 model_compression_toolkit/trainable_infrastructure/common/trainable_quantizer_config.py
 model_compression_toolkit/trainable_infrastructure/common/training_method.py
+model_compression_toolkit/trainable_infrastructure/common/util.py
 model_compression_toolkit/trainable_infrastructure/keras/__init__.py
+model_compression_toolkit/trainable_infrastructure/keras/annealing_schedulers.py
 model_compression_toolkit/trainable_infrastructure/keras/base_keras_quantizer.py
 model_compression_toolkit/trainable_infrastructure/keras/config_serialization.py
 model_compression_toolkit/trainable_infrastructure/keras/load_model.py
@@ -527,7 +529,6 @@ model_compression_toolkit/trainable_infrastructure/pytorch/__init__.py
 model_compression_toolkit/trainable_infrastructure/pytorch/annealing_schedulers.py
 model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py
 model_compression_toolkit/trainable_infrastructure/pytorch/quantizer_utils.py
-model_compression_toolkit/trainable_infrastructure/pytorch/util.py
 model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/__init__.py
 model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/base_activation_quantizer.py
 model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/lsq/__init__.py
@@ -566,6 +567,10 @@ tests_pytest/__init__.py
 tests_pytest/keras/__init__.py
 tests_pytest/keras/core/__init__.py
 tests_pytest/keras/core/test_data_util.py
+tests_pytest/keras/gptq/__init__.py
+tests_pytest/keras/gptq/test_gradual_act_quantization.py
+tests_pytest/keras/trainable_infrastructure/__init__.py
+tests_pytest/keras/trainable_infrastructure/test_linear_annealing.py
 tests_pytest/pytorch/__init__.py
 tests_pytest/pytorch/core/__init__.py
 tests_pytest/pytorch/core/test_data_util.py

{mct-nightly-2.2.0.20241026.508 → mct-nightly-2.2.0.20241027.532}/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__ = "2.2.0.20241026.000508"
+__version__ = "2.2.0.20241027.000532"

{mct-nightly-2.2.0.20241026.508 → mct-nightly-2.2.0.20241027.532}/model_compression_toolkit/gptq/common/gptq_constants.py

@@ -22,4 +22,11 @@ SOFT_ROUNDING_ZETA = 1.1
 
 # GPTQ config constant
 QUANT_PARAM_LEARNING_STR = 'quantization_parameter_learning'
-MAX_LSB_STR = 'max_lsbs_change_map'
+MAX_LSB_STR = 'max_lsbs_change_map'
+
+# GPTQ learning hyperparameters
+LR_DEFAULT = 3e-2
+LR_REST_DEFAULT = 1e-4
+LR_BIAS_DEFAULT = 1e-3
+LR_QUANTIZATION_PARAM_DEFAULT = 1e-3
+GPTQ_MOMENTUM = 0.9

{mct-nightly-2.2.0.20241026.508/model_compression_toolkit/gptq/pytorch/quantizer → mct-nightly-2.2.0.20241027.532/model_compression_toolkit/gptq/common}/gradual_activation_quantization.py

@@ -13,23 +13,23 @@
 # limitations under the License.
 # ==============================================================================
 from functools import partial
-from typing import Callable
+from typing import Callable, Any
 
 from model_compression_toolkit.gptq import GradientPTQConfig, QFractionLinearAnnealingConfig
-from model_compression_toolkit.trainable_infrastructure import BasePytorchTrainableQuantizer
-
-from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import LinearAnnealingScheduler
+from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import BaseTrainableQuantizer
 
 
 def get_gradual_activation_quantizer_wrapper_factory(gptq_config: GradientPTQConfig,
-                                                     get_total_grad_steps_fn: Callable[[], int]) \
-        -> Callable[[BasePytorchTrainableQuantizer], 'GradualActivationQuantizerWrapper']:
+                                                     get_total_grad_steps_fn: Callable[[], int],
+                                                     fw_linear_annealing_scheduler: type) \
+        -> Callable[[Any], 'GradualActivationQuantizerWrapper']:
     """
     Get a factory for 'GradualActivationQuantizerWrapper'.
 
     Args:
         gptq_config: GPTQ configuration.
         get_total_grad_steps_fn: a callable to obtain the total expected number of gradient steps.
+        fw_linear_annealing_scheduler: LinearAnnealingScheduler implementation of the framework (tf/pytorch).
 
     Returns:
         A factory function to build 'GradualActivationQuantizerWrapper' from Quantizer.
@@ -40,9 +40,9 @@ def get_gradual_activation_quantizer_wrapper_factory(gptq_config: GradientPTQCon
     annealing_cfg = gptq_config.gradual_activation_quantization_config.q_fraction_scheduler_policy
     if isinstance(annealing_cfg, QFractionLinearAnnealingConfig):
         t_end = annealing_cfg.end_step or get_total_grad_steps_fn()
-        factor_scheduler = LinearAnnealingScheduler(t_start=annealing_cfg.start_step, t_end=t_end,
-                                                    initial_val=annealing_cfg.initial_q_fraction,
-                                                    target_val=annealing_cfg.target_q_fraction)
+        factor_scheduler = fw_linear_annealing_scheduler(t_start=annealing_cfg.start_step, t_end=t_end,
+                                                         initial_val=annealing_cfg.initial_q_fraction,
+                                                         target_val=annealing_cfg.target_q_fraction)
     else:
         raise ValueError(f'Unknown annealing policy {annealing_cfg}')
 
@@ -64,7 +64,7 @@ class GradualActivationQuantizerWrapper:
         quantizer: quantizer to wrap.
         q_fraction_scheduler: a callable that accepts a gradient step and returns the corresponding quantized fraction.
     """
-    def __init__(self, quantizer: BasePytorchTrainableQuantizer, q_fraction_scheduler: Callable[[int], float]):
+    def __init__(self, quantizer: BaseTrainableQuantizer, q_fraction_scheduler: Callable[[int], float]):
         self.quantizer = quantizer
         self.q_fraction_scheduler = q_fraction_scheduler
         self.step_cnt = 0

{mct-nightly-2.2.0.20241026.508/model_compression_toolkit/gptq/pytorch/quantizer → mct-nightly-2.2.0.20241027.532/model_compression_toolkit/gptq/common}/regularization_factory.py

@@ -12,17 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Callable
 
-from model_compression_toolkit.gptq import RoundingType, GradientPTQConfig
-from model_compression_toolkit.gptq.pytorch.quantizer.soft_rounding.soft_quantizer_reg import \
-    SoftQuantizerRegularization
-from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import LinearAnnealingScheduler
+from tqdm import tqdm
+from typing import Callable, Type
 
+from model_compression_toolkit.gptq import RoundingType, GradientPTQConfig
 
+# Common warmup fraction
 WARMUP_STEP_FRACTION = 0.2
 
-def get_regularization(gptq_config: GradientPTQConfig, get_total_grad_steps_fn: Callable[[], int]) -> Callable:
+
+def get_regularization(gptq_config: GradientPTQConfig,
+                       get_total_grad_steps_fn: Callable[[], int],
+                       SoftQuantizerRegularizationFWClass: Type,
+                       LinearAnnealingSchedulerFWClass: Type) -> Callable:
     """
     Returns a function that computes the regularization term for GPTQ training based on the given
     rounding type in the GPTQ configuration.
@@ -30,15 +33,26 @@ def get_regularization(gptq_config: GradientPTQConfig, get_total_grad_steps_fn:
     Args:
         gptq_config: A GPTQ configuration.
         get_total_grad_steps_fn: a callable to obtain the total expected number of gradient steps.
+        SoftQuantizerRegularizationFWClass: The class to use for soft quantizer regularization (framework-specific).
+        LinearAnnealingSchedulerFWClass: The class to use for the annealing scheduler (framework-specific).
 
-    Returns: A function for computing the regularization. If there is no regularization function defined for the given
-        rounding type, then it returns a function that just returns 0.
-
+    Returns:
+        Callable: A function for computing the regularization. If there is no regularization function
+        defined for the given rounding type, then it returns a function that just returns 0.
     """
     if gptq_config.rounding_type == RoundingType.SoftQuantizer:
         total_gradient_steps = get_total_grad_steps_fn()
         t_start = int(WARMUP_STEP_FRACTION * total_gradient_steps)
-        scheduler = LinearAnnealingScheduler(t_start=t_start, t_end=total_gradient_steps, initial_val=20, target_val=2)
-        return SoftQuantizerRegularization(scheduler)
+
+        # Directly initializing the scheduler within the method
+        scheduler = LinearAnnealingSchedulerFWClass(
+            t_start=t_start,
+            t_end=total_gradient_steps,
+            initial_val=20,
+            target_val=2
+        )
+
+        # Return the framework-specific soft quantizer regularization
+        return SoftQuantizerRegularizationFWClass(scheduler)
     else:
         return lambda *args, **kwargs: 0

{mct-nightly-2.2.0.20241026.508 → mct-nightly-2.2.0.20241027.532}/model_compression_toolkit/gptq/keras/gptq_training.py

@@ -26,9 +26,14 @@ from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder
 from model_compression_toolkit.core.keras.data_util import data_gen_to_dataloader
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
+from model_compression_toolkit.gptq.common.gradual_activation_quantization import \
+    get_gradual_activation_quantizer_wrapper_factory
+from model_compression_toolkit.gptq.common.regularization_factory import get_regularization
 from model_compression_toolkit.gptq.keras.quantizer.quantization_builder import quantization_builder
 from model_compression_toolkit.logger import Logger
 from mct_quantizers import KerasActivationQuantizationHolder
+from model_compression_toolkit.trainable_infrastructure.common.util import get_total_grad_steps
+from model_compression_toolkit.trainable_infrastructure.keras.annealing_schedulers import KerasLinearAnnealingScheduler
 
 if version.parse(tf.__version__) >= version.parse("2.13"):
     from keras.src.engine.base_layer import TensorFlowOpLayer
@@ -41,13 +46,12 @@ from model_compression_toolkit.gptq.common.gptq_training import GPTQTrainer
 from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.gptq.keras.graph_info import get_weights_for_loss, get_gptq_trainable_parameters
-from model_compression_toolkit.gptq.keras.quantizer.regularization_factory import get_regularization
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 import numpy as np
 import copy
 from model_compression_toolkit.core.keras.constants import BIAS, USE_BIAS
-
+from model_compression_toolkit.gptq.keras.quantizer.soft_rounding.soft_quantizer_reg import SoftQuantizerRegularization
 
 class KerasGPTQTrainer(GPTQTrainer):
     """
@@ -78,6 +82,15 @@ class KerasGPTQTrainer(GPTQTrainer):
             hessian_info_service: HessianScoresService for fetching and computing Hessian's approximation scores.
 
         """
+
+        def _get_total_grad_steps():
+            return get_total_grad_steps(representative_data_gen) * gptq_config.n_epochs
+
+        # This must be set before the model building (as it is required for activation holder construction),
+        # which occurs in the base constructor.
+        self.gradual_act_quantizer_wrapper_factory = get_gradual_activation_quantizer_wrapper_factory(
+            gptq_config, _get_total_grad_steps, KerasLinearAnnealingScheduler)
+
         super().__init__(graph_float,
                          graph_quant,
                          gptq_config,
@@ -119,7 +132,10 @@ class KerasGPTQTrainer(GPTQTrainer):
 
         self.weights_for_average_loss = self._get_compare_points_loss_weights()
 
-        self.reg_func = get_regularization(self.gptq_config, representative_data_gen)
+        self.reg_func = get_regularization(self.gptq_config,
+                                           _get_total_grad_steps,
+                                           SoftQuantizerRegularization,
+                                           KerasLinearAnnealingScheduler)
 
     def _get_compare_points_loss_weights(self):
         """ Get compare points weights for the distillation loss. """
@@ -185,14 +201,13 @@ class KerasGPTQTrainer(GPTQTrainer):
         _, activation_quantizers = quantization_builder(n, self.gptq_config) # TODO: split quantizers building into two functions: for weights and activations
 
         # Holder by definition uses a single quantizer for the activation quantization
-        # thus we make sure this is the only possible case (unless it's a node with no activation
-        # quantization, which in this case has an empty list).
-        if len(activation_quantizers) == 1:
-            return KerasActivationQuantizationHolder(activation_quantizers[0])
-
-        Logger.critical(f"'KerasActivationQuantizationHolder' is designed to support a single quantizer, "
-                        f"but {len(activation_quantizers)} quantizers were found for node '{n}'. "
-                        f"Ensure only one quantizer is configured for each node's activation.")
+        # thus we make sure this is the only possible case.
+        if len(activation_quantizers) != 1:
+            Logger.critical(f"'KerasActivationQuantizationHolder' is designed to support a single quantizer, "
+                            f"but {len(activation_quantizers)} quantizers were found for node '{n}'. "
+                            f"Ensure only one quantizer is configured for each node's activation.")
+        quantizer = self.gradual_act_quantizer_wrapper_factory(activation_quantizers[0])
+        return KerasActivationQuantizationHolder(quantizer)
 
     def build_gptq_model(self) -> Tuple[Model, UserInformation]:
         """

{mct-nightly-2.2.0.20241026.508 → mct-nightly-2.2.0.20241027.532}/model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -14,17 +14,18 @@
 # ==============================================================================
 import copy
 
-from typing import Callable, Tuple
+from typing import Callable, Tuple, Union
 from packaging import version
 
-from model_compression_toolkit.core.common.quantization.quantize_graph_weights import quantize_graph_weights
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
-from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT
+from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT, LR_DEFAULT, LR_REST_DEFAULT, \
+    LR_BIAS_DEFAULT, GPTQ_MOMENTUM
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import TENSORFLOW, ACT_HESSIAN_DEFAULT_BATCH_SIZE
 from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.user_info import UserInformation
-from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig, GPTQHessianScoresConfig
+from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig, GPTQHessianScoresConfig, \
+    GradualActivationQuantizationConfig
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core import CoreConfig
@@ -32,13 +33,8 @@ from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.gptq.runner import gptq_runner
 from model_compression_toolkit.core.analyzer import analyzer_model_quantization
 from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import TargetPlatformCapabilities
-from model_compression_toolkit.metadata import get_versions_dict, create_model_metadata
+from model_compression_toolkit.metadata import create_model_metadata
 
-LR_DEFAULT = 0.15
-LR_REST_DEFAULT = 1e-4
-LR_BIAS_DEFAULT = 1e-4
-LR_QUANTIZATION_PARAM_DEFAULT = 1e-3
-GPTQ_MOMENTUM = 0.9
 
 if FOUND_TF:
     import tensorflow as tf
@@ -54,25 +50,25 @@ if FOUND_TF:
 
     # As from TF2.9 optimizers package is changed
     if version.parse(tf.__version__) < version.parse("2.9"):
-        from keras.optimizer_v2.optimizer_v2 import OptimizerV2
+        from keras.optimizer_v2.optimizer_v2 import OptimizerV2  # pragma: no cover
     elif version.parse(tf.__version__) < version.parse("2.12"):
-        from keras.optimizers.optimizer_v2.optimizer_v2 import OptimizerV2
+        from keras.optimizers.optimizer_v2.optimizer_v2 import OptimizerV2  # pragma: no cover
     else:
         from tensorflow.python.keras.optimizer_v2.optimizer_v2 import OptimizerV2
 
     DEFAULT_KERAS_TPC = get_target_platform_capabilities(TENSORFLOW, DEFAULT_TP_MODEL)
 
-
     def get_keras_gptq_config(n_epochs: int,
-                              optimizer: OptimizerV2 = tf.keras.optimizers.Adam(learning_rate=LR_DEFAULT),
-                              optimizer_rest: OptimizerV2 = tf.keras.optimizers.Adam(learning_rate=LR_REST_DEFAULT),
+                              optimizer: OptimizerV2 = None,
+                              optimizer_rest: OptimizerV2 = None,
                               loss: Callable = GPTQMultipleTensorsLoss(),
                               log_function: Callable = None,
                               use_hessian_based_weights: bool = True,
                               regularization_factor: float = REG_DEFAULT,
-                              hessian_batch_size: int = ACT_HESSIAN_DEFAULT_BATCH_SIZE) -> GradientPTQConfig:
+                              hessian_batch_size: int = ACT_HESSIAN_DEFAULT_BATCH_SIZE,
+                              gradual_activation_quantization: Union[bool, GradualActivationQuantizationConfig] = False) -> GradientPTQConfig:
         """
-        Create a GradientPTQConfigV2 instance for Keras models.
+        Create a GradientPTQConfig instance for Keras models.
 
         args:
             n_epochs (int): Number of epochs for running the representative dataset for fine-tuning.
@@ -83,9 +79,10 @@ if FOUND_TF:
             use_hessian_based_weights (bool): Whether to use Hessian-based weights for weighted average loss.
             regularization_factor (float): A floating point number that defines the regularization factor.
             hessian_batch_size (int): Batch size for Hessian computation in Hessian-based weights GPTQ.
+            gradual_activation_quantization (bool, GradualActivationQuantizationConfig): If False, GradualActivationQuantization is disabled. If True, GradualActivationQuantization is enabled with the default settings. GradualActivationQuantizationConfig object can be passed to use non-default settings.
 
         returns:
-            a GradientPTQConfigV2 object to use when fine-tuning the quantized model using gptq.
+            a GradientPTQConfig object to use when fine-tuning the quantized model using gptq.
 
         Examples:
 
@@ -94,7 +91,7 @@ if FOUND_TF:
             >>> import model_compression_toolkit as mct
             >>> import tensorflow as tf
 
-            Create a GradientPTQConfigV2 to run for 5 epochs:
+            Create a GradientPTQConfig to run for 5 epochs:
 
             >>> gptq_conf = mct.gptq.get_keras_gptq_config(n_epochs=5)
 
@@ -102,11 +99,24 @@ if FOUND_TF:
 
             >>> gptq_conf = mct.gptq.get_keras_gptq_config(n_epochs=3, optimizer=tf.keras.optimizers.Nadam())
 
-            The configuration can be passed to :func:`~model_compression_toolkit.keras_post_training_quantization` in order to quantize a keras model using gptq.
+            The configuration can be passed to :func:`~model_compression_toolkit.keras_gradient_post_training_quantization` in order to quantize a keras model using gptq.
+
 
         """
+        optimizer = optimizer or tf.keras.optimizers.Adam(learning_rate=LR_DEFAULT)
+        optimizer_rest = optimizer_rest or tf.keras.optimizers.Adam(learning_rate=LR_REST_DEFAULT)
+
         bias_optimizer = tf.keras.optimizers.SGD(learning_rate=LR_BIAS_DEFAULT,
                                                  momentum=GPTQ_MOMENTUM)
+
+        if isinstance(gradual_activation_quantization, bool):
+            gradual_quant_config = GradualActivationQuantizationConfig() if gradual_activation_quantization else None
+        elif isinstance(gradual_activation_quantization, GradualActivationQuantizationConfig):
+            gradual_quant_config = gradual_activation_quantization
+        else:
+            raise TypeError(f'gradual_activation_quantization argument should be bool or '
+                            f'GradualActivationQuantizationConfig, received {type(gradual_activation_quantization)}')
+
         return GradientPTQConfig(n_epochs,
                                  optimizer,
                                  optimizer_rest=optimizer_rest,
@@ -116,7 +126,8 @@ if FOUND_TF:
                                  optimizer_bias=bias_optimizer,
                                  use_hessian_based_weights=use_hessian_based_weights,
                                  regularization_factor=regularization_factor,
-                                 hessian_weights_config=GPTQHessianScoresConfig(hessian_batch_size=hessian_batch_size))
+                                 hessian_weights_config=GPTQHessianScoresConfig(hessian_batch_size=hessian_batch_size),
+                                 gradual_activation_quantization_config=gradual_quant_config)
 
 
     def keras_gradient_post_training_quantization(in_model: Model, representative_data_gen: Callable,
@@ -251,13 +262,13 @@ if FOUND_TF:
 else:
     # If tensorflow is not installed,
     # we raise an exception when trying to use these functions.
-    def get_keras_gptq_config(*args, **kwargs):
+    def get_keras_gptq_config(*args, **kwargs):  # pragma: no cover
         Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
                         "get_keras_gptq_config. The 'tensorflow' package is missing or is "
                         "installed with a version higher than 2.15.")  # pragma: no cover
 
 
-    def keras_gradient_post_training_quantization(*args, **kwargs):
+    def keras_gradient_post_training_quantization(*args, **kwargs):  # pragma: no cover
         Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
                         "keras_gradient_post_training_quantization. The 'tensorflow' package is missing or is "
-                        "installed with a version higher than 2.15.")  # pragma: no cover
+                        "installed with a version higher than 2.15.")

{mct-nightly-2.2.0.20241026.508 → mct-nightly-2.2.0.20241027.532}/model_compression_toolkit/gptq/keras/quantizer/quantization_builder.py

@@ -16,18 +16,18 @@ from typing import Dict, List, Tuple
 
 from model_compression_toolkit.gptq import GradientPTQConfig
 from model_compression_toolkit.core import common
-from model_compression_toolkit.exporter.model_wrapper.keras.builder.node_to_quantizer import \
-    get_inferable_quantizer_kwargs
 from model_compression_toolkit.gptq.keras.quantizer.base_keras_gptq_quantizer import BaseKerasGPTQTrainableQuantizer
 from mct_quantizers import QuantizationTarget
-from mct_quantizers.common.get_quantizers import get_inferable_quantizer_class
 from mct_quantizers.keras.quantizers import BaseKerasInferableQuantizer
 
 from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.trainable_infrastructure import TrainingMethod
 from model_compression_toolkit.trainable_infrastructure.common.get_quantizer_config import \
-    get_trainable_quantizer_weights_config
+    get_trainable_quantizer_weights_config, get_trainable_quantizer_activation_config
 from model_compression_toolkit.trainable_infrastructure.common.get_quantizers import \
     get_trainable_quantizer_class
+from model_compression_toolkit.trainable_infrastructure.keras.activation_quantizers.base_activation_quantizer import \
+    BaseKerasActivationTrainableQuantizer
 
 
 def quantization_builder(n: common.BaseNode,
@@ -70,12 +70,13 @@ def quantization_builder(n: common.BaseNode,
 
         quant_method = n.final_activation_quantization_cfg.activation_quantization_method
 
-        quantizer_class = get_inferable_quantizer_class(quant_target=QuantizationTarget.Activation,
+        quantizer_class = get_trainable_quantizer_class(quant_target=QuantizationTarget.Activation,
+                                                        quantizer_id=TrainingMethod.STE,
                                                         quant_method=quant_method,
-                                                        quantizer_base_class=BaseKerasInferableQuantizer)
+                                                        quantizer_base_class=BaseKerasActivationTrainableQuantizer)
+        cfg = get_trainable_quantizer_activation_config(n, None)
 
-        kwargs = get_inferable_quantizer_kwargs(n.final_activation_quantization_cfg, QuantizationTarget.Activation)
-
-        activation_quantizers.append(quantizer_class(**kwargs))
+        # freeze_quant_params is True since in GPTQ the activation quantization parameters should not be trained.
+        activation_quantizers.append(quantizer_class(cfg, freeze_quant_params=True))
 
     return weights_quantizers, activation_quantizers

{mct-nightly-2.2.0.20241026.508 → mct-nightly-2.2.0.20241027.532}/model_compression_toolkit/gptq/keras/quantizer/soft_rounding/soft_quantizer_reg.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import List
+from typing import List, Callable
 
 import tensorflow as tf
 from keras import Model
@@ -22,61 +22,21 @@ from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribut
 from model_compression_toolkit.trainable_infrastructure import KerasTrainableQuantizationWrapper
 
 
-class LinearTempDecay:
-    """
-    Annealing process for the soft quantizer regularization temperature term.
-    """
-
-    def __init__(self, t_max: int, rel_start_decay: float = 0.2, start_b: int = 20, end_b: int = 2):
-        """
-        Initializes a LinearTempDecay object.
-
-        Args:
-            t_max: maximal time step.
-            rel_start_decay: Decay step size at the beginning of the process.
-            start_b: Starting value of the regularization term.
-            end_b: Target value of the regularization term.
-        """
-
-        self.t_max = t_max
-        self.start_decay = rel_start_decay * t_max
-        self.start_b = start_b
-        self.end_b = end_b
-
-    def __call__(self, t: int) -> float:
-        """
-        Cosine annealing scheduler for soft quantizer regularization temperature term.
-
-        Args:
-            t: The current time step.
-
-        Returns: Scheduled temperature.
-        """
-
-        is_before_start_decay = tf.cast(t < self.start_decay, tf.float32)
-
-        rel_t = (t - self.start_decay) / (self.t_max - self.start_decay)
-
-        return self.start_b * is_before_start_decay + \
-               (1 - is_before_start_decay) * \
-               (self.end_b + (self.start_b - self.end_b) * tf.math.maximum(0.0, (1 - rel_t)))
-
 
 class SoftQuantizerRegularization:
     """
     A class to handle the computation of soft quantizer regularization for GPTQ training.
     """
 
-    def __init__(self, total_gradient_steps: int):
+    def __init__(self, beta_scheduler: Callable[[int], float]):
         """
         Initializes the regularization computation object with a LinearDecay object.
 
         Args:
-            total_gradient_steps: The number of gradient steps during optimization.
+            beta_scheduler: a callable that accepts current time step and returns a corresponding beta value.
         """
         # Initializing the temperature decay according to the number of expected gradient steps
-        self.linear_decay = LinearTempDecay(total_gradient_steps)
-
+        self.beta_scheduler = beta_scheduler
         self.count_iter = tf.Variable(0.)
 
 
@@ -91,7 +51,7 @@ class SoftQuantizerRegularization:
         Returns: Regularization value.
         """
         soft_reg_aux: List[tf.Tensor] = []
-        b = self.linear_decay(self.count_iter.value())
+        b = self.beta_scheduler(self.count_iter.value())
         for layer in model.layers:
             if isinstance(layer, KerasTrainableQuantizationWrapper):
                 kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=type(layer.layer),

{mct-nightly-2.2.0.20241026.508 → mct-nightly-2.2.0.20241027.532}/model_compression_toolkit/gptq/pytorch/gptq_training.py

@@ -17,15 +17,18 @@ from typing import Callable, List, Tuple, Union, Generator
 
 import numpy as np
 import torch
-from mct_quantizers import PytorchQuantizationWrapper, PytorchActivationQuantizationHolder
 from torch.nn import Module
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
+from model_compression_toolkit.gptq.common.gradual_activation_quantization import get_gradual_activation_quantizer_wrapper_factory
+from model_compression_toolkit.gptq.common.regularization_factory import get_regularization
+
 from model_compression_toolkit.core.common import Graph, BaseNode
 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.common.hessian import HessianInfoService, HessianScoresGranularity
+
 from model_compression_toolkit.core.pytorch.back2framework.pytorch_model_builder import PyTorchModelBuilder
 from model_compression_toolkit.core.pytorch.constants import BIAS
 from model_compression_toolkit.core.pytorch.data_util import FixedDatasetFromGenerator, IterableDatasetFromGenerator, \
@@ -34,14 +37,15 @@ from model_compression_toolkit.core.pytorch.utils import to_torch_tensor, set_mo
 from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
 from model_compression_toolkit.gptq.common.gptq_training import GPTQTrainer
-from model_compression_toolkit.gptq.pytorch.graph_info import get_gptq_trainable_parameters, \
-    get_weights_for_loss
-from model_compression_toolkit.gptq.pytorch.quantizer.gradual_activation_quantization import \
-    get_gradual_activation_quantizer_wrapper_factory
+from model_compression_toolkit.gptq.pytorch.graph_info import get_gptq_trainable_parameters, get_weights_for_loss
 from model_compression_toolkit.gptq.pytorch.quantizer.quantization_builder import quantization_builder
-from model_compression_toolkit.gptq.pytorch.quantizer.regularization_factory import get_regularization
+
+from mct_quantizers import PytorchQuantizationWrapper, PytorchActivationQuantizationHolder
+from model_compression_toolkit.trainable_infrastructure.common.util import get_total_grad_steps
+from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import PytorchLinearAnnealingScheduler
+from model_compression_toolkit.gptq.pytorch.quantizer.soft_rounding.soft_quantizer_reg import SoftQuantizerRegularization as PytorchSoftQuantizerRegularization
+
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.trainable_infrastructure.pytorch.util import get_total_grad_steps
 
 
 class PytorchGPTQTrainer(GPTQTrainer):
@@ -78,7 +82,7 @@ class PytorchGPTQTrainer(GPTQTrainer):
 
         # must be set prior to model building in the base class constructor
         self.gradual_act_quantizer_wrapper_factory = get_gradual_activation_quantizer_wrapper_factory(
-            gptq_config, _get_total_grad_steps)
+            gptq_config, _get_total_grad_steps, PytorchLinearAnnealingScheduler)
 
         super().__init__(graph_float,
                          graph_quant,
@@ -121,7 +125,7 @@ class PytorchGPTQTrainer(GPTQTrainer):
         else:
             self.train_dataloader = self._prepare_train_dataloader_for_non_sla(representative_data_gen)
 
-        self.reg_func = get_regularization(self.gptq_config, _get_total_grad_steps)
+        self.reg_func = get_regularization(self.gptq_config, _get_total_grad_steps, PytorchSoftQuantizerRegularization, PytorchLinearAnnealingScheduler)
 
     def _prepare_train_dataloader_sla(self, data_gen_fn: Callable[[], Generator]) -> DataLoader:
         """