mct-nightly 1.8.0.27022023.post430__py3-none-any.whl → 1.8.0.27032023.post403__py3-none-any.whl

model_compression_toolkit/gptq/common/gptq_config.py

@@ -16,9 +16,7 @@ from enum import Enum
 from typing import Callable, Any, Dict
 from model_compression_toolkit.core.common.defaultdict import DefaultDict
 from model_compression_toolkit.core import common
-from model_compression_toolkit.gptq.common.gptq_constants import N_BATCHES_STR, QUANT_PARAM_LEARNING_STR, N_EPOCHS_STR, \
-    MAX_LSB_STR
-from model_compression_toolkit.gptq.common.gptq_quantizer_config import GPTQQuantizerConfig, SoftQuantizerConfig
+from model_compression_toolkit.gptq.common.gptq_constants import QUANT_PARAM_LEARNING_STR, MAX_LSB_STR, REG_DEFAULT
 
 
 class RoundingType(Enum):
@@ -31,30 +29,53 @@ class RoundingType(Enum):
     SoftQuantizer = 1
 
 
+class GPTQHessianWeightsConfig:
+    """
+    Configuration to use for computing the Hessian-based weights for GPTQ loss metric.
+    """
+
+    def __init__(self,
+                 hessians_num_samples: int = 16,
+                 norm_weights: bool = True,
+                 log_norm: bool = True,
+                 scale_log_norm: bool = False,
+                 hessians_n_iter: int = 50):
+
+        """
+        Initialize a GPTQHessianWeightsConfig.
+        Args:
+            hessians_num_samples (int): Number of samples to use for computing the Hessian-based weights.
+            norm_weights (bool): Whether to normalize the returned weights (to get values between 0 and 1).
+            log_norm (bool): Whether to use log normalization to the GPTQ Hessian-based weights.
+            scale_log_norm (bool): Whether to scale the final vector of the Hessian weights.
+            hessians_n_iter (int): Number of random iterations to run Hessian approximation for GPTQ weights.
+        """
+
+        self.hessians_num_samples = hessians_num_samples
+        self.norm_weights = norm_weights
+        self.log_norm = log_norm
+        self.scale_log_norm = scale_log_norm
+        self.hessians_n_iter = hessians_n_iter
+
+
 class GradientPTQConfig:
     """
     Configuration to use for quantization with GradientPTQ (experimental).
     """
 
-    def __init__(self,
-                 n_iter: int,
+    def __init__(self, n_iter: int,
                  optimizer: Any,
                  optimizer_rest: Any = None,
                  loss: Callable = None,
                  log_function: Callable = None,
                  train_bias: bool = True,
-                 quantization_parameters_learning: bool = False,
                  rounding_type: RoundingType = RoundingType.SoftQuantizer,
-                 lsb_change_per_bit_width: dict = DefaultDict({}, lambda: 1),
-                 eps: float = 1e-6,
-                 use_jac_based_weights: bool = True,
-                 num_samples_for_loss: int = 16,
-                 norm_weights: bool = False,
+                 use_hessian_based_weights: bool = True,
                  optimizer_quantization_parameter: Any = None,
                  optimizer_bias: Any = None,
-                 log_norm: bool = True,
-                 weights_n_iter: int = 50,
-                 quantizer_config: GPTQQuantizerConfig = SoftQuantizerConfig()):
+                 regularization_factor: float = REG_DEFAULT,
+                 hessian_weights_config: GPTQHessianWeightsConfig = GPTQHessianWeightsConfig(),
+                 gptq_quantizer_params_override: Dict[str, Any] = None):
         """
         Initialize a GradientPTQConfig.
 
@@ -67,18 +88,13 @@ class GradientPTQConfig:
              accordingly. see example in multiple_tensors_mse_loss
             log_function (Callable): Function to log information about the GPTQ process.
             train_bias (bool): Whether to update the bias during the training or not.
-            quantization_parameters_learning (bool): Whether to update the quantization param during the training or not.
             rounding_type (RoundingType): An enum that defines the rounding type.
-            lsb_change_per_bit_width (dict): Whether to update the bias during the training or not.
-            eps (float): A floating point value for numeric stability.
-            use_jac_based_weights (bool): Whether to use jacobian-based weights for weighted average loss.
-            num_samples_for_loss (int): Number of samples to use for computing the jacobian-based weights.
-            norm_weights (bool): Whether to normalize the returned weights (to get values between 0 and 1).
+            use_hessian_based_weights (bool): Whether to use Hessian-based weights for weighted average loss.
             optimizer_quantization_parameter (Any): Optimizer to override the rest optimizer  for quantizer parameters.
-            optimizer_bias (Any): Optimizer to override the rest optimizerfor bias.
-            log_norm (bool): Whether to use log normalization to the GPTQ Jacobian-based weights.
-            weights_n_iter (int): Number of random iterations to run Jacobian approximation for GPTQ weights.
-            quantizer_config (GPTQQuantizerConfig): A class that contains the quantizer specific config.
+            optimizer_bias (Any): Optimizer to override the rest optimizer for bias.
+            regularization_factor (float): A floating point number that defines the regularization factor.
+            hessian_weights_config (GPTQHessianWeightsConfig): A configuration that include all necessary arguments to run a computation of Hessian weights for the GPTQ loss.
+            gptq_quantizer_params_override (dict): A dictionary of parameters to override in GPTQ quantizer instantiation. Defaults to None (no parameters).
 
         """
         self.n_iter = n_iter
@@ -88,69 +104,36 @@ class GradientPTQConfig:
         self.log_function = log_function
         self.train_bias = train_bias
 
-        if quantization_parameters_learning and rounding_type == RoundingType.STE:
-            common.Logger.error("Quantization parameters learning is not supported with STE rounding.")
-
-        self.quantization_parameters_learning = quantization_parameters_learning
         self.rounding_type = rounding_type
-        self.lsb_change_per_bit_width = lsb_change_per_bit_width
-        self.eps = eps
-        self.use_jac_based_weights = use_jac_based_weights
-        self.num_samples_for_loss = num_samples_for_loss
-        self.norm_weights = norm_weights
+        self.use_hessian_based_weights = use_hessian_based_weights
         self.optimizer_quantization_parameter = optimizer_quantization_parameter
         self.optimizer_bias = optimizer_bias
-        self.log_norm = log_norm
-        self.weights_n_iter = weights_n_iter
+        self.regularization_factor = regularization_factor
+        self.hessian_weights_config = hessian_weights_config
 
-        if self._verify_quantizer_config(quantizer_config, rounding_type):
-            self.quantizer_config = quantizer_config
-        else:
-            common.Logger.error(f"Quantizer config of type {type(quantizer_config)} "
-                                f"is not suitable for rounding type {rounding_type}")
-
-
-    def _verify_quantizer_config(self, quantizer_config, rounding_type) -> bool:
-        """
-        Verifies that the given quantizer config matches the given rounding type.
-
-        Args:
-            quantizer_config: A quantizer config.
-            rounding_type: A RoundingType.
-
-        Returns: True if the quantizer config matches the rounding type, False otherwise.
-
-        """
-        if rounding_type == RoundingType.SoftQuantizer:
-            return type(quantizer_config) == SoftQuantizerConfig
-
-        # Here, we compare type() and not isinstance to exclude instance equality because of inheritance
-        return type(quantizer_config) == GPTQQuantizerConfig
+        # Since the default quantizer is soft quantizer, we initialize the gptq_quantizer_params_override dictionary
+        # with its extended params
+        self.gptq_quantizer_params_override = {QUANT_PARAM_LEARNING_STR: False} \
+            if gptq_quantizer_params_override is None else gptq_quantizer_params_override
 
 
 class GradientPTQConfigV2(GradientPTQConfig):
     """
     Configuration to use for quantization with GradientPTQV2 (experimental).
     """
-    def __init__(self,
-                 n_epochs: int,
+    def __init__(self, n_epochs: int,
                  optimizer: Any,
                  optimizer_rest: Any = None,
                  loss: Callable = None,
                  log_function: Callable = None,
                  train_bias: bool = True,
-                 quantization_parameters_learning: bool = False,
                  rounding_type: RoundingType = RoundingType.SoftQuantizer,
-                 lsb_change_per_bit_width: dict = DefaultDict({}, lambda: 1),
-                 eps: float = 1e-6,
-                 use_jac_based_weights: bool = True,
-                 num_samples_for_loss: int = 16,
-                 norm_weights: bool = False,
+                 use_hessian_based_weights: bool = True,
                  optimizer_quantization_parameter: Any = None,
                  optimizer_bias: Any = None,
-                 log_norm: bool = True,
-                 weights_n_iter: int = 50,
-                 quantizer_config: GPTQQuantizerConfig = SoftQuantizerConfig()):
+                 regularization_factor: float = REG_DEFAULT,
+                 hessian_weights_config: GPTQHessianWeightsConfig = GPTQHessianWeightsConfig(),
+                 gptq_quantizer_params_override: Dict[str, Any] = None):
         """
         Initialize a GradientPTQConfigV2.
 
@@ -163,18 +146,13 @@ class GradientPTQConfigV2(GradientPTQConfig):
              accordingly. see example in multiple_tensors_mse_loss
             log_function (Callable): Function to log information about the GPTQ process.
             train_bias (bool): Whether to update the bias during the training or not.
-            quantization_parameters_learning (bool): Whether to update the quantization param during the training or not.
             rounding_type (RoundingType): An enum that defines the rounding type.
-            lsb_change_per_bit_width (dict): Whether to update the bias during the training or not.
-            eps (float): A floating point value for numeric stability.
-            use_jac_based_weights (bool): Whether to use jacobian-based weights for weighted average loss.
-            num_samples_for_loss (int): Number of samples to use for computing the jacobian-based weights.
-            norm_weights (bool): Whether to normalize the returned weights (to get values between 0 and 1).
+            use_hessian_based_weights (bool): Whether to use Hessian-based weights for weighted average loss.
             optimizer_quantization_parameter (Any): Optimizer to override the rest optimizer  for quantizer parameters.
             optimizer_bias (Any): Optimizer to override the rest optimizerfor bias.
-            log_norm (bool): Whether to use log normalization to the GPTQ Jacobian-based weights.
-            weights_n_iter (int): Number of random iterations to run Jacobian approximation for GPTQ weights.
-            quantizer_config (Any): A class that contains the quantizer specific config.
+            regularization_factor (float): A floating point number that defines the regularization factor.
+            hessian_weights_config (GPTQHessianWeightsConfig): A configuration that include all necessary arguments to run a computation of Hessian weights for the GPTQ loss.
+            gptq_quantizer_params_override (dict): A dictionary of parameters to override in GPTQ quantizer instantiation. Defaults to None (no parameters).
 
         """
 
@@ -184,18 +162,13 @@ class GradientPTQConfigV2(GradientPTQConfig):
                          loss=loss,
                          log_function=log_function,
                          train_bias=train_bias,
-                         quantization_parameters_learning=quantization_parameters_learning,
                          rounding_type=rounding_type,
-                         lsb_change_per_bit_width=lsb_change_per_bit_width,
-                         eps=eps,
-                         use_jac_based_weights=use_jac_based_weights,
-                         num_samples_for_loss=num_samples_for_loss,
-                         norm_weights=norm_weights,
+                         use_hessian_based_weights=use_hessian_based_weights,
                          optimizer_quantization_parameter=optimizer_quantization_parameter,
                          optimizer_bias=optimizer_bias,
-                         log_norm=log_norm,
-                         weights_n_iter=weights_n_iter,
-                         quantizer_config=quantizer_config)
+                         regularization_factor=regularization_factor,
+                         hessian_weights_config=hessian_weights_config,
+                         gptq_quantizer_params_override=gptq_quantizer_params_override)
         self.n_epochs = n_epochs
 
     @classmethod
@@ -212,22 +185,3 @@ class GradientPTQConfigV2(GradientPTQConfig):
         v1_params = config_v1.__dict__
         v1_params = {k: v for k, v in v1_params.items() if k != 'n_iter'}
         return cls(n_epochs, **v1_params)
-
-    def get_extended_quantizer_parametes(self) -> Dict[str, Any]:
-        """
-        Return a dictionary with a mapping to necessary additional parameters for initializing the GPTQ quantizer.
-
-        Returns: A dictionary with parameters for initializing a quantizer.
-
-        """
-
-        if self.rounding_type == RoundingType.SoftQuantizer:
-            return {N_BATCHES_STR: self.quantizer_config.n_batches,
-                    QUANT_PARAM_LEARNING_STR: self.quantization_parameters_learning,
-                    N_EPOCHS_STR: self.n_epochs}
-        elif self.rounding_type == RoundingType.STE:
-            return {MAX_LSB_STR: self.lsb_change_per_bit_width}
-
-        return {}
-
-

model_compression_toolkit/gptq/common/gptq_constants.py

@@ -2,7 +2,6 @@
 AUXVAR = 'auxvar_tensor'
 ITERVAR = 'iteration_variable'
 SCALE_TENSOR = "scale_ptq_tensor"
-GPTQ_ITER = "gptq_iter"
 AUXSHIFT = 'shift'
 WEIGHTS_QUANTIZATION_PARAMS = 'weights_quantization_params'
 PTQ_MIN_RANGE = "min_range"
@@ -11,22 +10,16 @@ PTQ_THRESHOLD = "ptq_threshold"
 SCALE_PTQ = "scale"
 
 # Default quantizer values
-N_EPOCHS = 10000
 N_CYCLES = 4
 MIM_TEMP = 0.5
 MAX_TEMP = 1.0
 REG_DEFAULT = 0.01
-MAX_ITERATIONS_DEFAULT = 10000
 MAX_LSB_CHANGE = 1
 
 # Soft rounding arguments values
 SOFT_ROUNDING_GAMMA = -0.1
 SOFT_ROUNDING_ZETA = 1.1
-SOFT_ROUNDING_BETA = 2 / 3
 
 # GPTQ config constant
-REGULARIZATION_VALUES = "regularization_values"
-N_BATCHES_STR = 'n_batches'
 QUANT_PARAM_LEARNING_STR = 'quantization_parameter_learning'
-N_EPOCHS_STR = 'n_epochs'
 MAX_LSB_STR = 'max_lsbs_change_map'

model_compression_toolkit/gptq/common/gptq_training.py

@@ -16,10 +16,11 @@ import copy
 from abc import ABC, abstractmethod
 import numpy as np
 from typing import Callable, List, Any
-from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig, RoundingType
+from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig
 from model_compression_toolkit.core.common import Graph, Logger, BaseNode
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
+from model_compression_toolkit.gptq.common.gptq_constants import QUANT_PARAM_LEARNING_STR
 from model_compression_toolkit.gptq.common.gptq_graph import get_compare_points
 from model_compression_toolkit.core.common.model_builder_mode import ModelBuilderMode
 
@@ -34,8 +35,7 @@ class GPTQTrainer(ABC):
                  graph_quant: Graph,
                  gptq_config: GradientPTQConfig,
                  fw_impl: FrameworkImplementation,
-                 fw_info: FrameworkInfo,
-                 representative_data_gen: Callable):
+                 fw_info: FrameworkInfo):
         """
         Build two models from a graph: A teacher network (float model) and a student network (quantized model).
         Use the dataset generator to pass images through the teacher and student networks to get intermediate
@@ -48,7 +48,6 @@ class GPTQTrainer(ABC):
             gptq_config: GradientPTQConfig with parameters about the tuning process.
             fw_impl: Framework implementation
             fw_info: Framework information
-            representative_data_gen: Dataset to use for inputs of the models.
         """
         self.graph_float = copy.deepcopy(graph_float)
         self.graph_quant = copy.deepcopy(graph_quant)
@@ -66,10 +65,6 @@ class GPTQTrainer(ABC):
                                                                        append2output=self.compare_points,
                                                                        fw_info=self.fw_info)
 
-        if self.gptq_config.rounding_type == RoundingType.SoftQuantizer:
-            # dry run on the representative dataset to count number of batches
-            self.count_num_batches_for_training(representative_data_gen)
-
         self.fxp_model, self.gptq_user_info = self.build_gptq_model()
 
     def get_optimizer_with_param(self,
@@ -88,8 +83,10 @@ class GPTQTrainer(ABC):
 
         w2train = [*flattened_trainable_weights]
 
+        quant_params_learning = self.gptq_config.gptq_quantizer_params_override.get(QUANT_PARAM_LEARNING_STR, False)
+
         optimizer_with_param = [(self.gptq_config.optimizer, w2train)]
-        if self.gptq_config.train_bias or self.gptq_config.quantization_parameters_learning:
+        if self.gptq_config.train_bias or quant_params_learning:
             w2train_res = []
             if self.gptq_config.train_bias:
                 if self.gptq_config.optimizer_bias is not None:
@@ -99,7 +96,7 @@ class GPTQTrainer(ABC):
                     if self.gptq_config.optimizer_rest is None:
                         Logger.error(  # pragma: no cover
                             "To enable bias micro training an additional optimizer is required, please define the optimizer_rest")
-            if self.gptq_config.quantization_parameters_learning:
+            if quant_params_learning:
                 if self.gptq_config.optimizer_quantization_parameter is not None:  # Ability to override optimizer
                     optimizer_with_param.append((self.gptq_config.optimizer_quantization_parameter,
                                                  trainable_quantization_parameters))
@@ -107,25 +104,32 @@ class GPTQTrainer(ABC):
                     w2train_res.extend(trainable_quantization_parameters)
                 if self.gptq_config.optimizer_rest is None:
                     Logger.error(  # pragma: no cover
-                        "To enable bias micro training an additional optimizer is required, please define the optimizer_rest")
-            optimizer_with_param.append((self.gptq_config.optimizer_rest, w2train_res))
+                        "To enable quantization parameters micro training an additional optimizer is required, please define the optimizer_rest")
+            if len(w2train_res) > 0:
+                # Either bias or quantization parameters are trainable but did not provide a specific optimizer,
+                # so we should use optimizer_rest to train them
+                if self.gptq_config.optimizer_rest is None:
+                    Logger.error(  # pragma: no cover
+                        "To enable bias or quantization parameters micro training an additional optimizer is required, please define the optimizer_rest")
+                optimizer_with_param.append((self.gptq_config.optimizer_rest, w2train_res))
 
         return optimizer_with_param
 
 
-    def compute_jacobian_based_weights(self,
-                                       representative_data_gen: Callable) -> np.ndarray:
+    def compute_hessian_based_weights(self,
+                                      representative_data_gen: Callable) -> np.ndarray:
         """
-        Computes the jacobian-based weights using the framework's model_grad method per batch of images.
+        Computes the Hessian-based weights using the framework's model_grad method per batch of images.
 
         Args:
-            representative_data_gen: Dataset used for inference to compute the jacobian-based weights.
+            representative_data_gen: Dataset used for inference to compute the Hessian-based weights.
 
         Returns: A vector of weights, one for each compare point,
         to be used for the loss metric weighted average computation when running GPTQ training.
         """
-        if self.gptq_config.use_jac_based_weights:
-            images = self._generate_images_batch(representative_data_gen, self.gptq_config.num_samples_for_loss)
+        if self.gptq_config.use_hessian_based_weights:
+            images = self._generate_images_batch(representative_data_gen,
+                                                 self.gptq_config.hessian_weights_config.hessians_num_samples)
 
             model_output_replacement = self._get_model_output_replacement()
 
@@ -143,17 +147,18 @@ class GPTQTrainer(ABC):
                                                              output_list=model_output_replacement,
                                                              all_outputs_indices=[],
                                                              alpha=0,
-                                                             norm_weights=self.gptq_config.norm_weights,
-                                                             n_iter=self.gptq_config.weights_n_iter)
+                                                             norm_weights=self.gptq_config.hessian_weights_config.norm_weights,
+                                                             n_iter=self.gptq_config.hessian_weights_config.hessians_n_iter)
                 points_apprx_jacobians_weights.append(image_ip_gradients)
-            if self.gptq_config.log_norm:
+            if self.gptq_config.hessian_weights_config.log_norm:
                 mean_jacobian_weights = np.mean(points_apprx_jacobians_weights, axis=0)
                 mean_jacobian_weights = np.where(mean_jacobian_weights != 0, mean_jacobian_weights,
                                                  np.partition(mean_jacobian_weights, 1)[1])
                 log_weights = np.log10(mean_jacobian_weights)
 
-                # To add scaling to the normalized weights replace return statement with the following line:
-                # return log_weights - np.min(log_weights) / (np.max(log_weights) - np.min(log_weights))
+                if self.gptq_config.hessian_weights_config.scale_log_norm:
+                    return (log_weights - np.min(log_weights)) / (np.max(log_weights) - np.min(log_weights))
+
                 return log_weights - np.min(log_weights)
             else:
                 return np.mean(points_apprx_jacobians_weights, axis=0)
@@ -249,21 +254,6 @@ class GPTQTrainer(ABC):
             replacement_outputs.append(prev_node)
         return replacement_outputs
 
-    def count_num_batches_for_training(self, representative_data_gen):
-        """
-        Runs a "dry-run" of the representative dataset to count the number of batches for each training epoch.
-
-        Args:
-            representative_data_gen: A callable method to retrieve images from Dataset.
-
-        Returns: The number of batches for each training epoch.
-
-        """
-        num_batches = 0
-        for _ in representative_data_gen():
-            num_batches += 1
-        self.gptq_config.quantizer_config.set_num_batches(num_batches)
-
 
 def gptq_training(graph_float: Graph,
                   graph_quant: Graph,

model_compression_toolkit/gptq/keras/gptq_training.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from functools import partial
 from typing import Callable, List, Tuple, Union
 
 import tensorflow as tf
@@ -23,7 +22,6 @@ from tqdm import tqdm
 # As from Tensorflow 2.6, keras is a separate package and some classes should be imported differently.
 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.gptq.common.gptq_constants import REGULARIZATION_VALUES
 from packaging import version
 
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
@@ -37,15 +35,15 @@ else:
 
 from model_compression_toolkit.core import common
 from model_compression_toolkit.gptq.common.gptq_training import GPTQTrainer
-from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfigV2, RoundingType
+from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfigV2
 from model_compression_toolkit.core.common import Graph
-from model_compression_toolkit.gptq.keras.graph_info import get_weights_for_loss, \
-    get_soft_rounding_reg, get_gptq_trainable_parameters
+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, KERNEL
+from model_compression_toolkit.core.keras.constants import BIAS, USE_BIAS
 from model_compression_toolkit import quantizers_infrastructure as qi
 
 
@@ -79,13 +77,12 @@ class KerasGPTQTrainer(GPTQTrainer):
                          graph_quant,
                          gptq_config,
                          fw_impl,
-                         fw_info,
-                         representative_data_gen)
+                         fw_info)
 
         self.loss_list = []
         self.input_scale = 1
 
-        trainable_weights, bias_weights, trainable_threshold, temperature_weights = get_gptq_trainable_parameters(
+        trainable_weights, bias_weights, trainable_threshold = get_gptq_trainable_parameters(
             self.fxp_model,
             fw_info,
             add_bias=gptq_config.train_bias)
@@ -112,7 +109,9 @@ class KerasGPTQTrainer(GPTQTrainer):
         else:
             self.input_scale = self.gptq_user_info.input_scale
 
-        self.weights_for_average_loss = self.compute_jacobian_based_weights(representative_data_gen)
+        self.weights_for_average_loss = self.compute_hessian_based_weights(representative_data_gen)
+
+        self.reg_func = get_regularization(self.gptq_config, representative_data_gen)
 
     def _is_gptq_applicable(self,
                             node: common.BaseNode) -> bool:
@@ -195,9 +194,7 @@ class KerasGPTQTrainer(GPTQTrainer):
                                                self.compare_points_std,
                                                self.weights_for_average_loss)
 
-            reg_value = self.gptq_config.quantizer_config.get_regularization_value(
-                self.fxp_model,
-                **{REGULARIZATION_VALUES: self._get_quantizer_regularization_values(self.gptq_config.rounding_type)})
+            reg_value = self.reg_func(self.fxp_model, self.gptq_config.regularization_factor)
 
             loss_value += reg_value
 
@@ -319,18 +316,3 @@ class KerasGPTQTrainer(GPTQTrainer):
                         node.set_weights_by_keys(BIAS, new_bias)
 
         return graph
-
-    def _get_quantizer_regularization_values(self, rounding_type: RoundingType) -> List[tf.Tensor]:
-        """
-        Mapping between a rounding type to its matching regularization method.
-
-        Args:
-            rounding_type: GPTQ rounding type.
-
-        Returns: A regularization computation method.
-
-        """
-        if rounding_type == RoundingType.SoftQuantizer:
-            return get_soft_rounding_reg(self.fxp_model)
-        else:
-            return []

model_compression_toolkit/gptq/keras/graph_info.py

@@ -13,23 +13,21 @@
 # limitations under the License.
 # ==============================================================================
 
-
 import tensorflow as tf
 from typing import Tuple, List
-
 from model_compression_toolkit.core.keras.constants import USE_BIAS
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from tensorflow.keras.models import Model
-
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
 from model_compression_toolkit.quantizers_infrastructure import KerasQuantizationWrapper
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
 def get_gptq_trainable_parameters(fxp_model: Model,
                                   fw_info: FrameworkInfo,
                                   add_bias: bool = False) -> (
-        List[tf.Variable], List[tf.Variable], List[tf.Variable], List[tf.Variable], List[tf.Variable]):
+        List[tf.Variable], List[tf.Variable], List[tf.Variable]):
     """
     Get trainable parameters from all layers in a model
 
@@ -45,16 +43,17 @@ def get_gptq_trainable_parameters(fxp_model: Model,
     trainable_weights: List[tf.Tensor] = []
     trainable_threshold: List[tf.Tensor] = []
     bias_weights: List[List[tf.Tensor]] = []
-    temperature_weights: List[tf.Tensor] = []
 
     for layer in fxp_model.layers:
         if isinstance(layer, KerasQuantizationWrapper):
             kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=type(layer.layer),
                                                                   fw_info=DEFAULT_KERAS_INFO)
 
-            # collect trainable weights per layer
-            layer_trainable_weights = layer.weights_quantizers[kernel_attribute].get_aux_variable()
-            layer_trainable_threshold = layer.weights_quantizers[kernel_attribute].get_quantization_variable()
+            # collect trainable weights per quantizer
+            quantizer_trainable_weights = layer.weights_quantizers[kernel_attribute].get_trainable_variables(VariableGroup.WEIGHTS)
+            quantizer_trainable_threshold = layer.weights_quantizers[kernel_attribute].get_trainable_variables(VariableGroup.QPARAMS)
+            trainable_weights.append(quantizer_trainable_weights)
+            trainable_threshold.extend(quantizer_trainable_threshold)
 
             if add_bias:
                 kernel_ops_attrs = fw_info.kernel_ops_attributes_mapping.get(type(layer.layer))
@@ -62,10 +61,8 @@ def get_gptq_trainable_parameters(fxp_model: Model,
                            and layer.layer.get_config().get(USE_BIAS)
                 if use_bias is not None and use_bias:
                     bias_weights.append([layer.layer.bias])
-            trainable_weights.append(layer_trainable_weights)
-            trainable_threshold.extend(layer_trainable_threshold)
 
-    return trainable_weights, bias_weights, trainable_threshold, temperature_weights
+    return trainable_weights, bias_weights, trainable_threshold
 
 
 def get_weights_for_loss(fxp_model: Model) -> Tuple[List[list], List[list]]:
@@ -95,25 +92,3 @@ def get_weights_for_loss(fxp_model: Model) -> Tuple[List[list], List[list]]:
             fxp_weights_list.append(_layer_fxp_weights)
 
     return flp_weights_list, fxp_weights_list
-
-
-# TODO: this function need to move to location that is relevant only for soft quantizer -
-#  once deciding how to handle GPTQ quantizers regularization.
-def get_soft_rounding_reg(fxp_model: Model) -> List[tf.Tensor]:
-    """
-    This function returns the soft quantizer regularization values for SoftRounding.
-
-    Args:
-        fxp_model: A model to be quantized with SoftRounding.
-
-    Returns: A list of tensors.
-    """
-
-    soft_reg_aux: List[tf.Tensor] = []
-    for layer in fxp_model.layers:
-        if isinstance(layer, KerasQuantizationWrapper):
-            kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=type(layer.layer),
-                                                                  fw_info=DEFAULT_KERAS_INFO)
-
-            soft_reg_aux.append(layer.weights_quantizers[kernel_attribute].get_regularization())
-    return soft_reg_aux

model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -85,24 +85,18 @@ if common.constants.FOUND_TF:
 
             Create a GradientPTQConfigV2 to run for 5 epochs:
 
-            >>> gptq_conf = mct.get_keras_gptq_config(n_epochs=5)
+            >>> gptq_conf = mct.gptq.get_keras_gptq_config(n_epochs=5)
 
             Other Tensorflow optimizers can be passed:
 
-            >>> gptq_conf = mct.get_keras_gptq_config(n_epochs=3, optimizer=tf.keras.optimizers.Nadam())
+            >>> 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.
 
         """
         bias_optimizer = tf.keras.optimizers.SGD(learning_rate=LR_BIAS_DEFAULT, momentum=GPTQ_MOMENTUM)
-        return GradientPTQConfigV2(n_epochs,
-                                   optimizer,
-                                   optimizer_rest=optimizer_rest,
-                                   loss=loss,
-                                   log_function=log_function,
-                                   train_bias=True,
-                                   quantization_parameters_learning=True,
-                                   optimizer_bias=bias_optimizer)
+        return GradientPTQConfigV2(n_epochs, optimizer, optimizer_rest=optimizer_rest, loss=loss,
+                                   log_function=log_function, train_bias=True, optimizer_bias=bias_optimizer)
 
 
     def keras_gradient_post_training_quantization_experimental(in_model: Model,
@@ -181,11 +175,11 @@ if common.constants.FOUND_TF:
 
             Create GPTQ config:
 
-            >>> gptq_config = mct.get_keras_gptq_config(n_epochs=1)
+            >>> 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.keras_gradient_post_training_quantization_experimental(model, repr_datagen, gptq_config, target_kpi=kpi, core_config=config)
+            >>> quantized_model, quantization_info = mct.gptq.keras_gradient_post_training_quantization_experimental(model, repr_datagen, gptq_config, target_kpi=kpi, core_config=config)
 
         """
         KerasModelValidation(model=in_model,