mct-nightly 1.8.0.1032023.post429__py3-none-any.whl → 1.8.0.1042023.post423__py3-none-any.whl

model_compression_toolkit/gptq/pytorch/graph_info.py

@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import torch
 import torch.nn as nn
 from typing import List
-
 from model_compression_toolkit.core.pytorch.constants import BIAS
 from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
 from model_compression_toolkit.quantizers_infrastructure import PytorchQuantizationWrapper
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
 def get_gptq_trainable_parameters(fxp_model: nn.Module,
@@ -38,24 +39,26 @@ def get_gptq_trainable_parameters(fxp_model: nn.Module,
     trainable_aux_weights = nn.ParameterList()
     trainable_threshold = nn.ParameterList()
     trainable_bias = nn.ParameterList()
-    trainable_temperature = nn.ParameterList()
 
     for layer in fxp_model.modules():
         if isinstance(layer, PytorchQuantizationWrapper):
             kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=type(layer.layer),
                                                                   fw_info=DEFAULT_PYTORCH_INFO)
 
-            trainable_aux_weights.extend(layer.weights_quantizers[kernel_attribute].get_aux_variable())
-            trainable_threshold.extend(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_aux_weights.extend(quantizer_trainable_weights)
+            trainable_threshold.extend(quantizer_trainable_threshold)
 
             if add_bias and hasattr(layer.layer, BIAS):
                 bias = getattr(layer.layer, BIAS)
                 trainable_bias.append(bias)
 
-    return trainable_aux_weights, trainable_bias, trainable_threshold, trainable_temperature
+    return trainable_aux_weights, trainable_bias, trainable_threshold
 
 
-def get_weights_for_loss(fxp_model: nn.Module) -> [List, List]:
+def get_weights_for_loss(fxp_model: nn.Module) -> [List[nn.Parameter], List[torch.Tensor]]:
     """
     Get all float and quantized kernels for the GPTQ loss
 

model_compression_toolkit/gptq/pytorch/quantization_facade.py

@@ -17,14 +17,15 @@ from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.constants import FOUND_TORCH
 from model_compression_toolkit.core.common import Logger
 from model_compression_toolkit.core.common.constants import PYTORCH
-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.target_platform import TargetPlatformCapabilities
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
 from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
+from model_compression_toolkit.gptq.keras.quantization_facade import GPTQ_MOMENTUM
 from model_compression_toolkit.gptq.runner import gptq_runner
 from model_compression_toolkit.core.exporter import export_model
 from model_compression_toolkit.core.analyzer import analyzer_model_quantization
-from model_compression_toolkit import CoreConfig, GPTQQuantizerConfig
+from model_compression_toolkit import CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
     MixedPrecisionQuantizationConfigV2
 
@@ -71,33 +72,19 @@ if FOUND_TORCH:
             Import MCT and Create a GradientPTQConfigV2 to run for 5 epochs:
 
             >>> import model_compression_toolkit as mct
-            >>> gptq_conf = mct.get_pytorch_gptq_config(n_epochs=5)
+            >>> gptq_conf = mct.gptq.get_pytorch_gptq_config(n_epochs=5)
 
             Other PyTorch optimizers can be passed with dummy params:
 
             >>> import torch
-            >>> gptq_conf = mct.get_pytorch_gptq_config(n_epochs=3, optimizer=torch.optim.Adam([torch.Tensor(1)]))
+            >>> gptq_conf = mct.gptq.get_pytorch_gptq_config(n_epochs=3, optimizer=torch.optim.Adam([torch.Tensor(1)]))
 
             The configuration can be passed to :func:`~model_compression_toolkit.pytorch_post_training_quantization` in order to quantize a pytorch model using gptq.
 
         """
-        bias_optimizer = Adam([torch.Tensor([])], lr=LR_BIAS_DEFAULT)
-        optimizer_quantization_parameter = Adam([torch.Tensor([])], lr=LR_QUANTIZATION_PARAM_DEFAULT)
-        # TODO: Once implementing Soft Quantizer for GPTQ in Pytorch:
-        #  - change default quantization_parameters_learning to True.
-        #  - remove explicit rounding_type and quantizer_config (and let it use the default GradientPTQConfig).
-        return GradientPTQConfigV2(n_epochs,
-                                   optimizer,
-                                   optimizer_rest=optimizer_rest,
-                                   loss=loss,
-                                   log_function=log_function,
-                                   train_bias=True,
-                                   optimizer_quantization_parameter=optimizer_quantization_parameter,
-                                   optimizer_bias=bias_optimizer,
-                                   rounding_type=RoundingType.STE,
-                                   quantizer_config=GPTQQuantizerConfig(),
-                                   quantization_parameters_learning=False,
-                                   )
+        bias_optimizer = torch.optim.SGD([torch.Tensor([])], lr=LR_BIAS_DEFAULT, momentum=GPTQ_MOMENTUM)
+        return GradientPTQConfigV2(n_epochs, optimizer, optimizer_rest=optimizer_rest, loss=loss,
+                                   log_function=log_function, train_bias=True, optimizer_bias=bias_optimizer)
 
 
     def pytorch_gradient_post_training_quantization_experimental(model: Module,
@@ -159,7 +146,7 @@ if FOUND_TORCH:
 
             Pass the module, the representative dataset generator and the configuration (optional) to get a quantized module
 
-            >>> quantized_module, quantization_info = mct.pytorch_gradient_post_training_quantization_experimental(module, repr_datagen, core_config=config, gptq_config=gptq_conf)
+            >>> quantized_module, quantization_info = mct.gptq.pytorch_gradient_post_training_quantization_experimental(module, repr_datagen, core_config=config, gptq_config=gptq_conf)
 
         """
 

model_compression_toolkit/gptq/pytorch/quantizer/__init__.py

@@ -13,4 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-import model_compression_toolkit.gptq.pytorch.quantizer.ste_rounding.symmetric_ste
+import model_compression_toolkit.gptq.pytorch.quantizer.ste_rounding.symmetric_ste
+import model_compression_toolkit.gptq.pytorch.quantizer.soft_rounding.symmetric_soft_quantizer
+import model_compression_toolkit.gptq.pytorch.quantizer.soft_rounding.uniform_soft_quantizer

model_compression_toolkit/gptq/pytorch/quantizer/base_pytorch_gptq_quantizer.py

@@ -71,26 +71,6 @@ if FOUND_TORCH:
 
             return weights, quant_config, {}
 
-        def get_aux_variable(self) -> List[Tensor]:
-            """
-            This function return a list with the quantizer's quantization auxiliary variables.
-
-            Returns: A list with the quantization auxiliary variables.
-
-            """
-
-            return []  # pragma: no cover
-
-        def get_quantization_variable(self) -> List[Tensor]:
-            """
-            This function return a list with the quantizer's quantization parameters variables.
-
-            Returns: A list with the quantization parameters.
-
-            """
-
-            return []  # pragma: no cover
-
         @abstractmethod
         def get_quant_config(self):
             """

model_compression_toolkit/gptq/pytorch/quantizer/quant_utils.py

@@ -30,11 +30,20 @@ def calculate_delta(max_tensor: torch.Tensor,
                     num_bits: int,
                     signed: bool) -> torch.Tensor:
     """
-    Compute the step size for the quantization.
+    Compute the step size for the symmetric quantization.
     """
     return max_tensor / (2 ** (num_bits - int(signed)))
 
 
+def calculate_delta_uniform(min_tensor: torch.Tensor,
+                            max_tensor: torch.Tensor,
+                            num_bits: int) -> torch.Tensor:
+    """
+    Compute the step size for the uniform quantization.
+    """
+    return (max_tensor-min_tensor) / (2 ** num_bits - 1)
+
+
 def ste_ceil(x: torch.Tensor) -> torch.Tensor:
     """
     Return the ceil values of a tensor.

model_compression_toolkit/gptq/pytorch/quantizer/quantization_builder.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 from typing import List, Dict, 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.pytorch.constants import KERNEL
 from model_compression_toolkit.exporter.model_wrapper.pytorch.builder.node_to_quantizer import \
@@ -59,7 +59,7 @@ def quantization_builder(n: common.BaseNode,
                                                         quant_method=quant_method,
                                                         quantizer_base_class=BasePytorchGPTQTrainableQuantizer)
         weights_quantizers.update({KERNEL: 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():
         quant_method = n.final_activation_quantization_cfg.activation_quantization_method

model_compression_toolkit/gptq/pytorch/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.pytorch.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/pytorch/quantizer/soft_rounding/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+# ==============================================================================

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

@@ -0,0 +1,115 @@
+# 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 torch
+import numpy as np
+from torch import nn
+
+from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
+from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
+from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
+from model_compression_toolkit.quantizers_infrastructure import PytorchQuantizationWrapper
+
+
+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: float) -> float:
+        """
+        Cosine annealing scheduler for soft quantizer regularization temperature term.
+
+        Args:
+            t: The current time step.
+
+        Returns: Scheduled temperature.
+        """
+
+        is_before_start_decay = (t < self.start_decay)
+
+        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) * torch.maximum(to_torch_tensor(np.array([0.0])),
+                                                                         to_torch_tensor(np.array((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: nn.Module, 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[torch.Tensor] = []
+        for layer in model.modules():
+            if isinstance(layer, PytorchQuantizationWrapper):
+                kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=type(layer.layer),
+                                                                      fw_info=DEFAULT_PYTORCH_INFO)
+
+                st = layer.weights_quantizers[kernel_attribute].get_soft_targets()
+                b = self.linear_decay(self.count_iter)
+
+                soft_reg_aux.append((1 - torch.pow(torch.abs(st - .5) * 2, b)).sum())
+
+        reg = 0
+
+        for sq in soft_reg_aux:
+            reg += sq
+
+        self.count_iter += 1
+
+        return entropy_reg * reg

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

@@ -0,0 +1,244 @@
+# 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.
+# ==============================================================================
+import torch
+import torch.nn as nn
+from typing import Dict
+import numpy as np
+
+from model_compression_toolkit.core.common import max_power_of_two
+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_config import RoundingType
+from model_compression_toolkit.gptq.pytorch.quantizer.base_pytorch_gptq_quantizer import \
+    BasePytorchGPTQTrainableQuantizer
+from model_compression_toolkit.core.pytorch.utils import to_torch_tensor, torch_tensor_to_numpy
+from model_compression_toolkit.gptq.pytorch.quantizer import quant_utils as qutils
+from model_compression_toolkit.gptq.common.gptq_constants import PTQ_THRESHOLD, SCALE_PTQ, \
+    SOFT_ROUNDING_GAMMA, SOFT_ROUNDING_ZETA, AUXVAR
+from model_compression_toolkit.core.common.constants import THRESHOLD, MIN_THRESHOLD
+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: torch.Tensor,
+                                      auxvar_tensor: torch.Tensor,
+                                      threshold_tensor: torch.Tensor,
+                                      num_bits: int,
+                                      signed: bool,
+                                      power_of_two: bool) -> torch.Tensor:
+    """
+    Quantize a tensor symmetrically for GPTQ quantizers.
+
+    Args:
+        input_tensor: Tensor to quantize. values of this tensor are not changed during gptq.
+        auxvar_tensor: Tensor that manifests the bit shift of the quantized weights due to gptq training.
+        threshold_tensor: Tensor with values to compute the threshold.
+        num_bits: Num of bits to use.
+        signed: Signedness of the quantization range.
+        power_of_two: Whether the threshold should be constrained or not.
+
+    Returns:
+        A quantized tensor.
+    """
+
+    if power_of_two:
+        threshold_tensor = qutils.power_of_two_max(threshold_tensor)
+    delta = qutils.calculate_delta(threshold_tensor, num_bits, signed)
+    with torch.no_grad():
+        input_tensor_int = torch.floor(input_tensor / delta)
+    tensor_q = input_tensor_int + auxvar_tensor
+    int_threshold = 2 ** (num_bits - int(signed))
+    return delta * qutils.ste_clip(tensor_q,
+                                   min_val=-int(signed) * int_threshold,
+                                   max_val=int_threshold - 1)
+
+
+@mark_quantizer(quantization_target=qi.QuantizationTarget.Weights,
+                quantization_method=[QuantizationMethod.POWER_OF_TWO, QuantizationMethod.SYMMETRIC],
+                quantizer_type=RoundingType.SoftQuantizer)
+class SymmetricSoftRoundingGPTQ(BasePytorchGPTQTrainableQuantizer):
+    """
+    Trainable symmetric quantizer to optimize the rounding of the quantized values using a soft quantization method.
+    """
+
+    def __init__(self,
+                 quantization_config: TrainableQuantizerWeightsConfig,
+                 quantization_parameter_learning: bool = False):
+        """
+        Construct a Pytorch model that utilize a fake weight quantizer of soft-quantizer for symmetric quantizer.
+
+        Args:
+            quantization_config: Trainable weights quantizer config.
+            quantization_parameter_learning (Bool): Whether to learn the threshold or not
+        """
+
+        super().__init__(quantization_config)
+        self.num_bits = quantization_config.weights_n_bits
+        self.per_channel = quantization_config.weights_per_channel_threshold
+
+        threshold_values = quantization_config.weights_quantization_params[THRESHOLD]
+        self.threshold_shape = np.asarray(threshold_values).shape
+        self.threshold_values = np.reshape(np.asarray(threshold_values), [-1]) if self.per_channel else float(
+            threshold_values)
+
+        self.quantization_axis = quantization_config.weights_channels_axis
+        self.power_of_two = quantization_config.weights_quantization_method == QuantizationMethod.POWER_OF_TWO
+        self.quantization_parameter_learning = quantization_parameter_learning
+
+        # gamma and zeta are stretch parameters for computing the rectified sigmoind function.
+        # See: https://arxiv.org/pdf/2004.10568.pdf
+        self.gamma = SOFT_ROUNDING_GAMMA
+        self.zeta = SOFT_ROUNDING_ZETA
+
+        self.quantizer_parameters = {}
+
+    def initialize_quantization(self,
+                                tensor_shape: torch.Size,
+                                name: str,
+                                layer: qi.PytorchQuantizationWrapper):
+        """
+        Add quantizer parameters to the quantizer parameters dictionary
+
+        Args:
+            tensor_shape: tensor shape of the quantized tensor.
+            name: Tensor name.
+            layer: Layer to quantize.
+        """
+
+        if self.per_channel:
+            threshold_tensor = to_torch_tensor(self.threshold_values)
+        else:
+            threshold_tensor = torch.tensor(self.threshold_values)
+        layer.register_parameter(f"{name}_{PTQ_THRESHOLD}",
+                                 nn.Parameter(threshold_tensor, requires_grad=False))
+
+        w = layer.layer.weight
+        delta = qutils.calculate_delta(threshold_tensor.reshape(self.threshold_shape), self.num_bits, signed=True)
+        w_clipped_normed = torch.clip(w / delta, -2**(self.num_bits-1), 2**(self.num_bits-1)-1)
+        rest = w_clipped_normed - torch.floor(w_clipped_normed)  # rest of rounding [0, 1)
+        # Note that (rest - self.gamma) can't be zero since rest is positive and gamma is negative, so the division
+        # is safe
+        alpha = -torch.log((self.zeta - self.gamma) / (rest - self.gamma) - 1)  # => sigmoid(alpha) = rest
+
+        layer.register_parameter(f"{name}_{AUXVAR}", nn.Parameter(alpha, requires_grad=True))
+
+        # save the quantizer added parameters for later calculations
+        self.add_quantizer_variable(PTQ_THRESHOLD, layer.get_parameter(f"{name}_{PTQ_THRESHOLD}"), VariableGroup.QPARAMS)
+        self.add_quantizer_variable(AUXVAR, layer.get_parameter(f"{name}_{AUXVAR}"), VariableGroup.WEIGHTS)
+
+        if self.quantization_parameter_learning:
+            if self.per_channel:
+                layer.register_parameter(f"{name}_{SCALE_PTQ}",
+                                         nn.Parameter(to_torch_tensor(torch.ones_like(torch.Tensor(self.threshold_values))),
+                                                      requires_grad=True))
+            else:
+                layer.register_parameter(f"{name}_{SCALE_PTQ}",
+                                         nn.Parameter(to_torch_tensor((torch.tensor([1.0], requires_grad=True)))))
+            self.add_quantizer_variable(SCALE_PTQ, layer.get_parameter(f"{name}_{SCALE_PTQ}"), VariableGroup.QPARAMS)
+
+    def get_soft_targets(self) -> torch.Tensor:
+        """
+        Computes the rectified sigmoid function for the quantization target parameters.
+
+        Returns:
+            A tensor with the soft rounding targets values.
+
+        """
+        scaled_sigmoid = torch.sigmoid(self.get_quantizer_variable(AUXVAR)) * (self.zeta - self.gamma) + self.gamma
+        return torch.clip(scaled_sigmoid, min=0, max=1)
+
+    def get_quant_config(self) -> Dict[str, np.ndarray]:
+        """
+        Returns the config used to edit NodeQuantizationConfig after GPTQ retraining
+
+        Returns:
+            A dictionary of attributes the quantize_config retraining has changed during GPTQ retraining.
+            Keys must match NodeQuantizationConfig attributes
+
+        """
+        old_threshold = torch_tensor_to_numpy(self.get_quantizer_variable(PTQ_THRESHOLD))
+        old_threshold = np.resize(old_threshold, self.threshold_shape)
+        if self.power_of_two:
+            old_threshold = max_power_of_two(old_threshold, MIN_THRESHOLD)
+        else:
+            if self.quantization_parameter_learning:
+                scale = torch.reshape(self.get_quantizer_variable(SCALE_PTQ), self.threshold_shape)
+                old_threshold = old_threshold * torch_tensor_to_numpy(scale)
+        old_threshold = old_threshold.reshape(self.threshold_shape)
+        return {THRESHOLD: old_threshold}
+
+    def __call__(self,
+                 inputs: nn.Parameter,
+                 training: bool) -> torch.Tensor:
+        """
+        Quantize a tensor.
+
+        Args:
+            inputs: Input tensor to quantize.
+            training: whether in training mode or not
+
+        Returns:
+            quantized tensor
+        """
+        auxvar = self.get_quantizer_variable(AUXVAR)
+        ptq_threshold_tensor = self.get_quantizer_variable(PTQ_THRESHOLD)
+
+        #####################################################
+        # Soft Rounding
+        #####################################################
+        aux_var = self.get_soft_targets()
+        if not training:
+            aux_var = (aux_var >= 0.5).to(auxvar.dtype)
+
+        if self.per_channel:
+            reshape_shape = get_threshold_reshape_shape(inputs.shape,
+                                                        quant_axis=self.quantization_axis,
+                                                        quant_axis_dim=-1)
+
+            ##########################################################
+            # Calculate soft rounding targets and optimized threshold
+            ##########################################################
+            ptq_threshold_tensor_hat = torch.reshape(ptq_threshold_tensor, reshape_shape)
+
+            #####################################################
+            # Quantized Input
+            #####################################################
+            q_tensor = soft_rounding_symmetric_quantizer(input_tensor=inputs,
+                                                         auxvar_tensor=aux_var,
+                                                         threshold_tensor=ptq_threshold_tensor_hat,
+                                                         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 = torch.reshape(self.get_quantizer_variable(SCALE_PTQ), reshape_shape)
+                q_tensor *= scale
+
+        else:
+            q_tensor = soft_rounding_symmetric_quantizer(input_tensor=inputs,
+                                                         auxvar_tensor=aux_var,
+                                                         threshold_tensor=ptq_threshold_tensor,
+                                                         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