mct-nightly 1.8.0.2032023.post428__py3-none-any.whl → 1.8.0.2042023.post413__py3-none-any.whl

model_compression_toolkit/gptq/keras/quantizer/quantization_builder.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 from typing import Dict, List, Tuple
 
-from model_compression_toolkit import GradientPTQConfigV2
+from model_compression_toolkit.gptq import GradientPTQConfigV2
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.exporter.model_wrapper.keras.builder.node_to_quantizer import \
@@ -61,7 +61,7 @@ def quantization_builder(n: common.BaseNode,
                                                               fw_info=DEFAULT_KERAS_INFO)
 
         weights_quantizers.update({kernel_attribute: quantizer_class(get_trainable_quantizer_weights_config(n),
-                                                                     **gptq_config.get_extended_quantizer_parametes())})
+                                                                     **gptq_config.gptq_quantizer_params_override)})
 
     activation_quantizers = []
     if n.is_activation_quantization_enabled():

model_compression_toolkit/gptq/keras/quantizer/regularization_factory.py

@@ -0,0 +1,45 @@
+# Copyright 2023 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from typing import Callable
+
+from model_compression_toolkit.gptq import RoundingType, GradientPTQConfigV2, GradientPTQConfig
+from model_compression_toolkit.gptq.keras.quantizer.soft_rounding.soft_quantizer_reg import \
+    SoftQuantizerRegularization
+
+
+def get_regularization(gptq_config: GradientPTQConfig, representative_data_gen: Callable) -> Callable:
+    """
+    Returns a function that computes the regularization term for GPTQ training based on the given
+    rounding type in the GPTQ configuration.
+
+    Args:
+        gptq_config: A GPTQ configuration.
+        representative_data_gen: Dataset used for the GPTQ training.
+
+    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.
+
+    """
+    if gptq_config.rounding_type == RoundingType.SoftQuantizer:
+        # dry run on the representative dataset to count number of batches
+        num_batches = 0
+        for _ in representative_data_gen():
+            num_batches += 1
+
+        n_epochs = GradientPTQConfigV2.from_v1(n_ptq_iter=num_batches, config_v1=gptq_config).n_epochs if \
+            not type(gptq_config) == GradientPTQConfigV2 else gptq_config.n_epochs
+        return SoftQuantizerRegularization(total_gradient_steps=num_batches * n_epochs)
+    else:
+        return lambda m, e_reg: 0

model_compression_toolkit/gptq/keras/quantizer/soft_rounding/soft_quantizer_reg.py

@@ -0,0 +1,112 @@
+# Copyright 2023 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from typing import List
+
+import tensorflow as tf
+from keras 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
+
+
+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):
+        """
+        Initializes the regularization computation object with a LinearDecay object.
+
+        Args:
+            total_gradient_steps: The number of gradient steps during optimization.
+        """
+        # Initializing the temperature decay according to the number of expected gradient steps
+        self.linear_decay = LinearTempDecay(total_gradient_steps)
+
+        self.count_iter = 0
+
+
+    def __call__(self, model: Model, entropy_reg: float):
+        """
+        Returns the soft quantizer regularization value for SoftRounding.
+
+        Args:
+            model: A model to be quantized with SoftRounding.
+            entropy_reg: Entropy value to scale the quantizer regularization.
+
+        Returns: Regularization value.
+        """
+
+        soft_reg_aux: List[tf.Tensor] = []
+        for layer in model.layers:
+            if isinstance(layer, KerasQuantizationWrapper):
+                kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=type(layer.layer),
+                                                                      fw_info=DEFAULT_KERAS_INFO)
+
+                st = layer.weights_quantizers[kernel_attribute].get_soft_targets()
+                b = self.linear_decay(self.count_iter)
+
+                soft_reg_aux.append(tf.reduce_sum(1 - tf.pow(tf.math.abs(st - .5) * 2, b)))
+
+        reg = 0
+
+        for sq in soft_reg_aux:
+            reg += sq
+
+        self.count_iter += 1
+
+        return entropy_reg * reg

model_compression_toolkit/gptq/keras/quantizer/soft_rounding/symmetric_soft_quantizer.py

@@ -16,22 +16,22 @@
 import tensorflow as tf
 import numpy as np
 
-from model_compression_toolkit import RoundingType
+from model_compression_toolkit.gptq import RoundingType
 from model_compression_toolkit import quantizers_infrastructure as qi
 from model_compression_toolkit.core.common import max_power_of_two
 from model_compression_toolkit.core.common.target_platform import QuantizationMethod
-from model_compression_toolkit.gptq.common.gptq_constants import PTQ_THRESHOLD, SCALE_PTQ, N_EPOCHS, \
-    MAX_ITERATIONS_DEFAULT, SOFT_ROUNDING_GAMMA, SOFT_ROUNDING_ZETA, SOFT_ROUNDING_BETA, GPTQ_ITER, AUXVAR
+from model_compression_toolkit.gptq.common.gptq_constants import PTQ_THRESHOLD, SCALE_PTQ, \
+    SOFT_ROUNDING_GAMMA, SOFT_ROUNDING_ZETA, AUXVAR
 from model_compression_toolkit.gptq.keras.quantizer import quant_utils as qutils
-from typing import Dict, Any, List
+from typing import Dict, Any
 from model_compression_toolkit.core.common.constants import THRESHOLD, MIN_THRESHOLD
-from model_compression_toolkit.core.common.logger import Logger
 from model_compression_toolkit.gptq.keras.quantizer.base_keras_gptq_quantizer import BaseKerasGPTQTrainableQuantizer
 from model_compression_toolkit.gptq.keras.quantizer.quant_utils import power_of_two_max, clip, calculate_delta
 from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.base_inferable_quantizer import mark_quantizer
 from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.quant_utils import \
     get_threshold_reshape_shape
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
 def soft_rounding_symmetric_quantizer(input_tensor: tf.Tensor,
@@ -66,46 +66,6 @@ def soft_rounding_symmetric_quantizer(input_tensor: tf.Tensor,
     return delta * clip(tensor_q, max_val=max_int, min_val=min_int)
 
 
-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)))
-
-
 @mark_quantizer(quantization_target=qi.QuantizationTarget.Weights,
                 quantization_method=[QuantizationMethod.POWER_OF_TWO, QuantizationMethod.SYMMETRIC],
                 quantizer_type=RoundingType.SoftQuantizer)
@@ -116,23 +76,15 @@ class SymmetricSoftRoundingGPTQ(BaseKerasGPTQTrainableQuantizer):
 
     def __init__(self,
                  quantization_config: TrainableQuantizerWeightsConfig,
-                 n_batches: int = None,
-                 quantization_parameter_learning: bool = False,
-                 n_epochs: int = N_EPOCHS):
+                 quantization_parameter_learning: bool = False):
         """
         Initialize a SymmetricSoftRoundingGPTQ object with parameters to use
         for the quantization.
 
         Args:
             quantization_config: Trainable weights quantizer config.
-            n_batches: The expected number of batches for each training epoch.
             quantization_parameter_learning: Whether to train the quantization threshold.
-            n_epochs: Number of epochs to run training for.
         """
-
-        if n_batches is None:
-            Logger.error("SymmetricSoftRoundingGPTQ got an uninitialized n_batches argument.")
-
         super().__init__(quantization_config)
         self.num_bits = quantization_config.weights_n_bits
         self.per_channel = quantization_config.weights_per_channel_threshold
@@ -148,36 +100,23 @@ class SymmetricSoftRoundingGPTQ(BaseKerasGPTQTrainableQuantizer):
         self.num_channels = len(self.threshold_values) if self.per_channel else 1
 
         # gamma and zeta are stretch parameters for computing the rectified sigmoind function.
-        # beta is used to set the regularization term.
         # See: https://arxiv.org/pdf/2004.10568.pdf
         self.gamma = SOFT_ROUNDING_GAMMA
         self.zeta = SOFT_ROUNDING_ZETA
-        self.beta = SOFT_ROUNDING_BETA
 
         self.quantizer_parameters = {}
 
-        # Initializing the temperature decay according to the number of expected gradient steps
-        if n_batches is None:
-            Logger.warning(f"Number of batches is not set correctly for the Soft Quantizer. A default value of "  # pragma: no cover
-                           f"{MAX_ITERATIONS_DEFAULT} is used to set the temperature decay which may affect the results.")
-
-        init_decay = MAX_ITERATIONS_DEFAULT if n_batches is None else n_epochs * n_batches
-        self.linear_decay = LinearTempDecay(init_decay)
-
     def initialize_quantization(self,
                                 tensor_shape: Any,
                                 name: str,
-                                layer: Any) -> Dict[Any, Any]:
+                                layer: Any):
         """
-        Return a dictionary of quantizer parameters and their names.
+        Add quantizer parameters to the quantizer parameters dictionary
 
         Args:
             tensor_shape: tensor shape of the quantized tensor.
             name: Tensor name.
             layer: Layer to quantize.
-
-        Returns:
-            Dictionary of parameters names to the variables.
         """
 
         if self.per_channel:
@@ -187,12 +126,6 @@ class SymmetricSoftRoundingGPTQ(BaseKerasGPTQTrainableQuantizer):
         else:
             reshape_shape = [self.num_channels]
 
-        ar_iter = layer.add_weight(
-            f"{name}_{GPTQ_ITER}",
-            shape=(),
-            initializer=tf.keras.initializers.Constant(0.0),
-            trainable=False)
-
         ptq_threshold_tensor = layer.add_weight(
             f"{name}_{PTQ_THRESHOLD}",
             shape=reshape_shape,
@@ -216,44 +149,17 @@ class SymmetricSoftRoundingGPTQ(BaseKerasGPTQTrainableQuantizer):
 
         auxvar_tensor.assign(alpha)
 
-        self.quantizer_parameters.update({AUXVAR: auxvar_tensor,
-                                          PTQ_THRESHOLD: ptq_threshold_tensor,
-                                          GPTQ_ITER: ar_iter})
+        # Add quantization variables
+        self.add_quantizer_variable(AUXVAR, auxvar_tensor, VariableGroup.WEIGHTS)
+        self.add_quantizer_variable(PTQ_THRESHOLD, ptq_threshold_tensor, VariableGroup.QPARAMS)
 
-        if self.quantization_parameter_learning:
+        if self.quantization_parameter_learning and not self.power_of_two:
             scale = layer.add_weight(
                 f"{name}_{SCALE_PTQ}",
                 shape=self.num_channels,
                 initializer=tf.keras.initializers.Constant(1.0),
                 trainable=True)
-            self.quantizer_parameters.update({SCALE_PTQ: scale})
-
-        return self.quantizer_parameters
-
-    def get_quantization_variable(self) -> List[tf.Tensor]:
-        """
-        This function return a list with the quantizer's quantization parameters variables.
-
-        Returns: A list with the quantization parameters if there are defined parameters.
-
-        """
-
-        if self.quantization_parameter_learning and not self.power_of_two:
-            return [self.quantizer_parameters[SCALE_PTQ]]
-        else:
-            return []
-
-    def get_regularization(self) -> tf.Tensor:
-        """
-        Computes the regularization term for the soft rounding loss.
-
-        Returns:
-            regularization term.
-        """
-
-        st = self.get_soft_targets()
-        b = self.linear_decay(self.ar_iter.value())
-        return tf.reduce_sum(1 - tf.pow(tf.math.abs(st - .5) * 2, b))
+            self.add_quantizer_variable(SCALE_PTQ, scale, VariableGroup.QPARAMS)
 
     def get_soft_targets(self) -> tf.Tensor:
         """
@@ -264,16 +170,7 @@ class SymmetricSoftRoundingGPTQ(BaseKerasGPTQTrainableQuantizer):
 
         """
         return qutils.clip(
-            tf.sigmoid(self.quantizer_parameters[AUXVAR]) * (self.zeta - self.gamma) + self.gamma, 1, 0)
-
-    def get_aux_variable(self) -> List[tf.Tensor]:
-        """
-        This function return a list with the quantizer's quantization auxiliary variables.
-
-        Returns: A list with the quantization auxiliary variables.
-
-        """
-        return [self.quantizer_parameters[AUXVAR]]
+            tf.sigmoid(self.get_quantizer_variable(AUXVAR)) * (self.zeta - self.gamma) + self.gamma, 1, 0)
 
     def __call__(self,
                  inputs: tf.Tensor,
@@ -289,8 +186,14 @@ class SymmetricSoftRoundingGPTQ(BaseKerasGPTQTrainableQuantizer):
             The quantized tensor.
         """
 
-        self.ar_iter = self.quantizer_parameters[GPTQ_ITER]
-        ptq_threshold_tensor = self.quantizer_parameters[PTQ_THRESHOLD]
+        ptq_threshold_tensor = self.get_quantizer_variable(PTQ_THRESHOLD)
+
+        #####################################################
+        # Soft Rounding
+        #####################################################
+        aux_var = self.get_soft_targets()
+        if not training:
+            aux_var = tf.cast(tf.math.greater_equal(aux_var, 0.5), tf.float32)
 
         if self.per_channel:
             reshape_shape = get_threshold_reshape_shape(inputs.shape,
@@ -301,15 +204,6 @@ class SymmetricSoftRoundingGPTQ(BaseKerasGPTQTrainableQuantizer):
             # Calculate soft rounding targets and optimized threshold
             ##########################################################
             ptq_threshold_tensor_hat = tf.reshape(ptq_threshold_tensor, reshape_shape)
-            aux_var = self.get_soft_targets()
-
-            #####################################################
-            # Soft Rounding
-            #####################################################
-            if training:
-                self.ar_iter.assign_add(1.0)
-            else:
-                aux_var = tf.cast(self.quantizer_parameters[AUXVAR] >= 0, tf.float32)
 
             #####################################################
             # Quantized Input
@@ -322,17 +216,22 @@ class SymmetricSoftRoundingGPTQ(BaseKerasGPTQTrainableQuantizer):
                                                          power_of_two=self.power_of_two)
 
             if self.quantization_parameter_learning and not self.power_of_two:
-                scale = tf.reshape(self.quantizer_parameters[SCALE_PTQ], reshape_shape)
+                scale = tf.reshape(self.get_quantizer_variable(SCALE_PTQ), reshape_shape)
                 q_tensor *= scale
 
-            return q_tensor
         else:
-            return soft_rounding_symmetric_quantizer(input_tensor=inputs,
-                                                     auxvar_tensor=self.quantizer_parameters[AUXVAR],
-                                                     threshold_tensor=ptq_threshold_tensor.value(),
-                                                     num_bits=self.num_bits,
-                                                     signed=True,
-                                                     power_of_two=self.power_of_two)
+            q_tensor = soft_rounding_symmetric_quantizer(input_tensor=inputs,
+                                                         auxvar_tensor=aux_var,
+                                                         threshold_tensor=ptq_threshold_tensor.value(),
+                                                         num_bits=self.num_bits,
+                                                         signed=True,
+                                                         power_of_two=self.power_of_two)
+
+            if self.quantization_parameter_learning and not self.power_of_two:
+                scale = self.get_quantizer_variable(SCALE_PTQ)
+                q_tensor *= scale
+
+        return q_tensor
 
     def get_quant_config(self) -> Dict[str, np.ndarray]:
         """
@@ -344,13 +243,13 @@ class SymmetricSoftRoundingGPTQ(BaseKerasGPTQTrainableQuantizer):
         """
 
         if self.power_of_two:
-            old_threshold = self.quantizer_parameters[PTQ_THRESHOLD]
+            old_threshold = self.get_quantizer_variable(PTQ_THRESHOLD)
             old_threshold = max_power_of_two(old_threshold, MIN_THRESHOLD)
 
         else:
-            old_threshold = self.quantizer_parameters[PTQ_THRESHOLD]
+            old_threshold = self.get_quantizer_variable(PTQ_THRESHOLD)
             if self.quantization_parameter_learning:
-                scale = tf.reshape(self.quantizer_parameters[SCALE_PTQ], self.threshold_shape)
+                scale = tf.reshape(self.get_quantizer_variable(SCALE_PTQ), self.threshold_shape)
                 old_threshold = old_threshold * scale
             old_threshold = old_threshold.numpy()
         old_threshold = old_threshold.reshape(self.threshold_shape)

model_compression_toolkit/gptq/keras/quantizer/ste_rounding/symmetric_ste.py

@@ -13,15 +13,15 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Dict, Any, List
+from typing import Dict, Any
 
 import numpy as np
 import tensorflow as tf
 
-from model_compression_toolkit import RoundingType
+from model_compression_toolkit.gptq import RoundingType
 from model_compression_toolkit import quantizers_infrastructure as qi
 from model_compression_toolkit.core.common.target_platform import QuantizationMethod
-from model_compression_toolkit.gptq.common.gptq_constants import GPTQ_ITER, AUXVAR, PTQ_THRESHOLD
+from model_compression_toolkit.gptq.common.gptq_constants import AUXVAR, PTQ_THRESHOLD
 from model_compression_toolkit.gptq.keras.quantizer import quant_utils as qutils
 from model_compression_toolkit.core.common.constants import THRESHOLD
 from model_compression_toolkit.core.common.defaultdict import DefaultDict
@@ -30,6 +30,7 @@ from model_compression_toolkit.quantizers_infrastructure import TrainableQuantiz
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.base_inferable_quantizer import mark_quantizer
 from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.quant_utils import \
     get_threshold_reshape_shape
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
 def pertubation_symmetric_quantizer(input_tensor: tf.Tensor,
@@ -96,30 +97,20 @@ class STEWeightGPTQQuantizer(BaseKerasGPTQTrainableQuantizer):
         self.quantization_axis = quantization_config.weights_channels_axis
         self.power_of_two = quantization_config.weights_quantization_method == QuantizationMethod.POWER_OF_TWO
         self.max_lsbs_change = max_lsbs_change_map.get(self.num_bits)
-        self.quantizer_parameters = {}
 
     def initialize_quantization(self,
                                 tensor_shape: Any,
                                 name: str,
-                                layer: Any) -> Dict[Any, Any]:
+                                layer: Any):
         """
-        Return a dictionary of quantizer parameters and their names.
+        Add quantizer parameters to the quantizer parameters dictionary
 
         Args:
             tensor_shape: tensor shape of the quantized tensor.
             name: Tensor name.
             layer: Layer to quantize.
-
-        Returns:
-            Dictionary of parameters names to the variables.
         """
 
-        ar_iter = layer.add_weight(
-            f"{name}_{GPTQ_ITER}",
-            shape=(),
-            initializer=tf.keras.initializers.Constant(0.0),
-            trainable=False)
-
         ptq_threshold_tensor = layer.add_weight(
             f"{name}_{PTQ_THRESHOLD}",
             shape=len(self.threshold_values) if self.per_channel else (),
@@ -135,10 +126,8 @@ class STEWeightGPTQQuantizer(BaseKerasGPTQTrainableQuantizer):
             trainable=True)
 
         # save the quantizer added parameters for later calculations
-        self.quantizer_parameters = {PTQ_THRESHOLD: ptq_threshold_tensor,
-                                     AUXVAR: auxvar_tensor,
-                                     GPTQ_ITER: ar_iter}
-        return self.quantizer_parameters
+        self.add_quantizer_variable(PTQ_THRESHOLD, ptq_threshold_tensor, VariableGroup.QPARAMS)
+        self.add_quantizer_variable(AUXVAR, auxvar_tensor, VariableGroup.WEIGHTS)
 
     def __call__(self,
                  inputs: tf.Tensor,
@@ -154,8 +143,8 @@ class STEWeightGPTQQuantizer(BaseKerasGPTQTrainableQuantizer):
             The quantized tensor.
         """
 
-        auxvar = self.quantizer_parameters[AUXVAR]
-        ptq_threshold_tensor = self.quantizer_parameters[PTQ_THRESHOLD]
+        auxvar = self.get_quantizer_variable(AUXVAR)
+        ptq_threshold_tensor = self.get_quantizer_variable(PTQ_THRESHOLD)
 
         if self.per_channel:
             reshape_shape = get_threshold_reshape_shape(inputs.shape,
@@ -178,25 +167,6 @@ class STEWeightGPTQQuantizer(BaseKerasGPTQTrainableQuantizer):
                                                    signed=True,
                                                    power_of_two=self.power_of_two)
 
-    def get_aux_variable(self) -> List[tf.Tensor]:
-        """
-        This function return a list with the quantizer's quantization auxiliary variables.
-
-        Returns: A list with the quantization auxiliary variables.
-
-        """
-
-        return [self.quantizer_parameters[AUXVAR]]
-
-    def get_quantization_variable(self) -> List[tf.Tensor]:
-        """
-        This function return a list with the quantizer's quantization parameters variables.
-
-        Returns: A list with the quantization parameters.
-
-        """
-
-        return [self.quantizer_parameters[PTQ_THRESHOLD]]
 
     def get_quant_config(self) -> Dict[str, np.ndarray]:
         """
@@ -207,5 +177,5 @@ class STEWeightGPTQQuantizer(BaseKerasGPTQTrainableQuantizer):
             Keys must match NodeQuantizationConfig attributes
 
         """
-        old_threshold = self.quantizer_parameters[PTQ_THRESHOLD]
+        old_threshold = self.get_quantizer_variable(PTQ_THRESHOLD)
         return {THRESHOLD: old_threshold.numpy().reshape(self.threshold_shape)}

model_compression_toolkit/gptq/pytorch/gptq_training.py

@@ -29,9 +29,11 @@ from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.pytorch.constants import BIAS
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor, set_model, torch_tensor_to_numpy
-from model_compression_toolkit.gptq.pytorch.graph_info import get_gptq_trainable_parameters, get_weights_for_loss
+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 import quantizers_infrastructure as qi
+from model_compression_toolkit.gptq.pytorch.quantizer.regularization_factory import get_regularization
 from model_compression_toolkit.quantizers_infrastructure import PytorchQuantizationWrapper
 
 
@@ -61,7 +63,7 @@ class PytorchGPTQTrainer(GPTQTrainer):
             fw_info: Framework information
             representative_data_gen: Dataset to use for inputs of the models.
         """
-        super().__init__(graph_float, graph_quant, gptq_config, fw_impl, fw_info, representative_data_gen)
+        super().__init__(graph_float, graph_quant, gptq_config, fw_impl, fw_info)
         self.loss_list = []
         self.input_scale = 1
         if self.float_user_info.input_scale != self.gptq_user_info.input_scale:
@@ -69,7 +71,7 @@ class PytorchGPTQTrainer(GPTQTrainer):
         else:
             self.input_scale = self.gptq_user_info.input_scale
 
-        trainable_weights, trainable_bias, trainable_threshold, trainable_temperature = get_gptq_trainable_parameters(
+        trainable_weights, trainable_bias, trainable_threshold = get_gptq_trainable_parameters(
             self.fxp_model,
             add_bias=self.gptq_config.train_bias)
 
@@ -84,7 +86,9 @@ class PytorchGPTQTrainer(GPTQTrainer):
                                                                   trainable_bias,
                                                                   trainable_threshold)
 
-        self.weights_for_average_loss = to_torch_tensor(self.compute_jacobian_based_weights(representative_data_gen))
+        self.weights_for_average_loss = to_torch_tensor(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: BaseNode) -> bool:
@@ -182,6 +186,10 @@ class PytorchGPTQTrainer(GPTQTrainer):
                                            self.compare_points_std,
                                            self.weights_for_average_loss)
 
+        reg_value = self.reg_func(self.fxp_model, self.gptq_config.regularization_factor)
+
+        loss_value += reg_value
+
         # Back-pass
         loss_value.backward()