mct-nightly 1.8.0.8032023.post421__py3-none-any.whl → 1.8.0.8052023.post414__py3-none-any.whl

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

@@ -13,23 +13,24 @@
 # 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.target_platform_capabilities.target_platform import QuantizationMethod
+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.constants import THRESHOLD
 from model_compression_toolkit.core.common.defaultdict import DefaultDict
 from model_compression_toolkit.gptq.keras.quantizer.base_keras_gptq_quantizer import BaseKerasGPTQTrainableQuantizer
 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 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_pytorch_implementation.py

@@ -0,0 +1,29 @@
+# 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 Type
+
+from model_compression_toolkit.core.pytorch.pytorch_implementation import PytorchImplementation
+from model_compression_toolkit.gptq.common.gptq_framework_implementation import GPTQFrameworkImplemantation
+from model_compression_toolkit.gptq.pytorch.gptq_training import PytorchGPTQTrainer
+
+
+class GPTQPytorchImplemantation(GPTQFrameworkImplemantation, PytorchImplementation):
+
+    def get_gptq_trainer_obj(self) -> Type[PytorchGPTQTrainer]:
+        """
+        Returns:  Pytorch object of GPTQTrainer
+        """
+        return PytorchGPTQTrainer

model_compression_toolkit/gptq/pytorch/gptq_training.py

@@ -19,21 +19,21 @@ from torch.nn import Module
 from tqdm import tqdm
 import copy
 import torch
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.pytorch.back2framework.pytorch_model_builder import PyTorchModelBuilder
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
 from model_compression_toolkit.gptq.common.gptq_training import GPTQTrainer
-from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfigV2, RoundingType
+from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfigV2
 from model_compression_toolkit.core.common import Graph, 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.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, get_soft_rounding_reg
+    get_weights_for_loss
 from model_compression_toolkit.gptq.pytorch.quantizer.quantization_builder import quantization_builder
-from model_compression_toolkit.gptq.common.gptq_constants import REGULARIZATION_VALUES
 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
 
 
@@ -63,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:
@@ -71,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)
 
@@ -86,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:
@@ -184,9 +186,7 @@ class PytorchGPTQTrainer(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
 
@@ -272,18 +272,3 @@ class PytorchGPTQTrainer(GPTQTrainer):
                 if hasattr(layer.layer, BIAS):
                     bias = getattr(layer.layer, BIAS)
                     bias.requires_grad = self.gptq_config.train_bias
-
-    def _get_quantizer_regularization_values(self, rounding_type: RoundingType) -> List[torch.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/pytorch/graph_info.py

@@ -15,11 +15,11 @@
 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,
@@ -39,21 +39,23 @@ 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[nn.Parameter], List[torch.Tensor]]:
@@ -77,25 +79,3 @@ def get_weights_for_loss(fxp_model: nn.Module) -> [List[nn.Parameter], List[torc
                 fxp_weights_list.append(quantizer(training=False, inputs=quantizer_vars))
 
     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: nn.Module) -> List[torch.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[torch.Tensor] = []
-    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)
-
-            soft_reg_aux.append(layer.weights_quantizers[kernel_attribute].get_regularization())
-    return soft_reg_aux

model_compression_toolkit/gptq/pytorch/quantization_facade.py

@@ -14,17 +14,18 @@
 # ==============================================================================
 from typing import Callable
 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.core.common.target_platform import TargetPlatformCapabilities
+from model_compression_toolkit.constants import FOUND_TORCH
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import PYTORCH
+from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfigV2
+from model_compression_toolkit.target_platform_capabilities.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.core import CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
     MixedPrecisionQuantizationConfigV2
 
@@ -35,8 +36,8 @@ LR_QUANTIZATION_PARAM_DEFAULT = 1e-4
 
 if FOUND_TORCH:
     from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
-    from model_compression_toolkit.core.pytorch.pytorch_implementation import PytorchImplementation
-    from model_compression_toolkit.core.pytorch.constants import DEFAULT_TP_MODEL
+    from model_compression_toolkit.gptq.pytorch.gptq_pytorch_implementation import GPTQPytorchImplemantation
+    from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from model_compression_toolkit.gptq.pytorch.gptq_loss import multiple_tensors_mse_loss
     from model_compression_toolkit.exporter.model_wrapper.pytorch.builder.fully_quantized_model_builder import get_exportable_pytorch_model
     import torch
@@ -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,
@@ -131,7 +118,7 @@ if FOUND_TORCH:
             core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
             gptq_config (GradientPTQConfigV2): Configuration for using gptq (e.g. optimizer).
             gptq_representative_data_gen (Callable): Dataset used for GPTQ training. If None defaults to representative_data_gen
-            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the PyTorch model according to. `Default PyTorch TPC <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/core/tpc_models/pytorch_tp_models/pytorch_default.py>`_
+            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the PyTorch model according to.
             new_experimental_exporter (bool): Whether exporting the quantized model using new exporter or not (in progress. Avoiding it for now is recommended).
 
         Returns:
@@ -155,26 +142,26 @@ if FOUND_TORCH:
 
             Create MCT core configurations with number of calibration iterations set to 1:
 
-            >>> config = mct.CoreConfig()
+            >>> config = mct.core.CoreConfig()
 
             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)
 
         """
 
         if core_config.mixed_precision_enable:
             if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfigV2):
-                common.Logger.error("Given quantization config to mixed-precision facade is not of type "
+                Logger.error("Given quantization config to mixed-precision facade is not of type "
                                     "MixedPrecisionQuantizationConfigV2. Please use keras_post_training_quantization "
                                     "API, or pass a valid mixed precision configuration.")  # pragma: no cover
 
-            common.Logger.info("Using experimental mixed-precision quantization. "
+            Logger.info("Using experimental mixed-precision quantization. "
                                "If you encounter an issue please file a bug.")
 
         tb_w = _init_tensorboard_writer(DEFAULT_PYTORCH_INFO)
 
-        fw_impl = PytorchImplementation()
+        fw_impl = GPTQPytorchImplemantation()
 
         # ---------------------- #
         # Core Runner
@@ -205,7 +192,7 @@ if FOUND_TORCH:
             Logger.warning('Using new experimental exported models. '
                            'Please do not use unless you are familiar with what you are doing')
 
-            return get_fully_quantized_pytorch_model(graph_gptq)
+            return get_exportable_pytorch_model(graph_gptq)
 
         return export_model(graph_gptq,
                             DEFAULT_PYTORCH_INFO,

model_compression_toolkit/gptq/pytorch/quantizer/__init__.py

@@ -15,3 +15,4 @@
 
 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

@@ -13,10 +13,10 @@
 # limitations under the License.
 # ==============================================================================
 from abc import abstractmethod
-from typing import Union, Dict, List
+from typing import Union, Dict
 
-from model_compression_toolkit.core.common.logger import Logger
-from model_compression_toolkit.core.common.constants import FOUND_TORCH
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import FOUND_TORCH
 from model_compression_toolkit.gptq.common.gptq_constants import WEIGHTS_QUANTIZATION_PARAMS
 
 from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig, \
@@ -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

@@ -14,9 +14,7 @@
 # ==============================================================================
 from typing import Union, Tuple
 import torch
-from torch.nn.functional import softmax, log_softmax, one_hot
-from model_compression_toolkit.core.common.constants import MIN_THRESHOLD
-from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
+from model_compression_toolkit.constants import MIN_THRESHOLD
 
 
 def power_of_two_max(max_tensor: torch.Tensor) -> torch.Tensor:
@@ -30,11 +28,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.
@@ -42,6 +49,13 @@ def ste_ceil(x: torch.Tensor) -> torch.Tensor:
     return (torch.ceil(x) - x).detach() + x
 
 
+def ste_floor(x: torch.Tensor) -> torch.Tensor:
+    """
+    Return the floor values of a tensor.
+    """
+    return (torch.floor(x) - x).detach() + x
+
+
 def ste_round(x: torch.Tensor) -> torch.Tensor:
     """
     Calculate the rounded 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