mct-nightly 1.8.0.4032023.post406__py3-none-any.whl → 1.8.0.4042023.post409__py3-none-any.whl

model_compression_toolkit/qat/pytorch/quantizer/ste_rounding/symmetric_ste.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Dict, Union
+from typing import Union
 
 import numpy as np
 import torch
@@ -31,6 +31,7 @@ from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructur
     ActivationSymmetricInferableQuantizer
 from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.trainable_quantizer_config import \
     TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
 @mark_quantizer(quantization_target=qi.QuantizationTarget.Weights,
@@ -66,22 +67,19 @@ class STEWeightQATQuantizer(BasePytorchQATTrainableQuantizer):
         self.max_int = (2 ** n_pos_bits) - 1
         self.min = delta * self.min_int
         self.max = delta * self.max_int
-        self.quantizer_parameters = {}
+
 
     def initialize_quantization(self,
                                 tensor_shape: torch.Size,
                                 name: str,
-                                layer: qi.PytorchQuantizationWrapper) -> Dict[str, nn.Parameter]:
+                                layer: qi.PytorchQuantizationWrapper):
         """
-        Add min and max variables to layer.
-        Args:
-            tensor_shape: Tensor shape the quantizer quantize.
-            name: Prefix of variables names.
-            layer: Layer to add the variables to. The variables are saved
-            in the layer's scope.
+        Add quantizer parameters to the quantizer parameters dictionary
 
-        Returns:
-            Dictionary of new variables.
+        Args:
+            tensor_shape: tensor shape of the quantized tensor.
+            name: Tensor name.
+            layer: Layer to quantize.
         """
 
         # Add threshold variables to layer.
@@ -89,9 +87,8 @@ class STEWeightQATQuantizer(BasePytorchQATTrainableQuantizer):
                                                                              requires_grad=False))
 
         # save the quantizer added parameters for later calculations
-        self.quantizer_parameters = {THRESHOLD_TENSOR: layer.get_parameter(name + "_" + THRESHOLD_TENSOR)}
+        self.add_quantizer_variable(THRESHOLD_TENSOR, layer.get_parameter(name + "_" + THRESHOLD_TENSOR), VariableGroup.QPARAMS)
 
-        return self.quantizer_parameters
 
     def __call__(self,
                  inputs: nn.Parameter,
@@ -116,7 +113,7 @@ class STEWeightQATQuantizer(BasePytorchQATTrainableQuantizer):
         Returns:
             A pytorch inferable quanizer object.
         """
-        np_threshold = self.quantizer_parameters[THRESHOLD_TENSOR].cpu().detach().numpy().flatten()
+        np_threshold = self.get_quantizer_variable(THRESHOLD_TENSOR).cpu().detach().numpy().flatten()
         if self.power_of_two:
             pot_threshold = 2 ** np.ceil(np.log2(np_threshold))
             return WeightsPOTInferableQuantizer(num_bits=self.num_bits,
@@ -153,20 +150,23 @@ class STEActivationQATQuantizer(BasePytorchQATTrainableQuantizer):
         np_threshold_values = quantization_config.activation_quantization_params[C.THRESHOLD]
         self.threshold_tensor = torch.Tensor([np_threshold_values])
         self.num_bits = quantization_config.activation_n_bits
-        self.quantizer_parameters = {}
 
     def initialize_quantization(self,
                                 tensor_shape: torch.Size,
                                 name: str,
-                                layer: qi.PytorchQuantizationWrapper) -> Dict[str, nn.Parameter]:
+                                layer: qi.PytorchQuantizationWrapper):
         """
-        Add threshold variables to layer.
+        Add quantizer parameters to the quantizer parameters dictionary
+
+        Args:
+            tensor_shape: tensor shape of the quantized tensor.
+            name: Tensor name.
+            layer: Layer to quantize.
         """
         layer.register_parameter(name, nn.Parameter(to_torch_tensor(self.threshold_tensor), requires_grad=True))
 
         # save the quantizer added parameters for later calculations
-        self.quantizer_parameters = {THRESHOLD_TENSOR: layer.get_parameter(name)}
-        return self.quantizer_parameters
+        self.add_quantizer_variable(THRESHOLD_TENSOR, layer.get_parameter(name), VariableGroup.QPARAMS)
 
     def __call__(self,
                  inputs: torch.Tensor,
@@ -181,7 +181,7 @@ class STEActivationQATQuantizer(BasePytorchQATTrainableQuantizer):
             The quantized tensor.
         """
 
-        _t = self.quantizer_parameters[THRESHOLD_TENSOR]
+        _t = self.get_quantizer_variable(THRESHOLD_TENSOR)
         q_tensor = symmetric_quantizer(inputs, _t, self.num_bits, sign=self.sign)
         return q_tensor
 
@@ -192,7 +192,7 @@ class STEActivationQATQuantizer(BasePytorchQATTrainableQuantizer):
         Returns:
             A pytorch inferable quanizer object.
         """
-        np_threshold = self.quantizer_parameters[THRESHOLD_TENSOR].cpu().detach().numpy()
+        np_threshold = self.get_quantizer_variable(THRESHOLD_TENSOR).cpu().detach().numpy()
         if self.power_of_two:
             pot_threshold = np.power(2.0, np.ceil(np.log2(np_threshold)))
             return ActivationPOTInferableQuantizer(num_bits=self.num_bits,

model_compression_toolkit/qat/pytorch/quantizer/ste_rounding/uniform_ste.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Dict
-
 import numpy as np
 import torch
 import torch.nn as nn
@@ -32,6 +30,7 @@ from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructur
     WeightsUniformInferableQuantizer, ActivationUniformInferableQuantizer
 from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.trainable_quantizer_config import \
     TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
 @mark_quantizer(quantization_target=qi.QuantizationTarget.Weights,
@@ -64,22 +63,18 @@ class STEUniformWeightQATQuantizer(BasePytorchQATTrainableQuantizer):
                               [-1]) if self.quantization_config.weights_per_channel_threshold else float(
             self.min_values)
 
-        self.quantizer_parameters = {}
 
     def initialize_quantization(self,
                                 tensor_shape: torch.Size,
                                 name: str,
-                                layer: qi.PytorchQuantizationWrapper) -> Dict[str, nn.Parameter]:
+                                layer: qi.PytorchQuantizationWrapper):
         """
-        Add min and max variables to layer.
-        Args:
-            tensor_shape: Tensor shape the quantizer quantize.
-            name: Prefix of variables names.
-            layer: Layer to add the variables to. The variables are saved
-            in the layer's scope.
+        Add quantizer parameters to the quantizer parameters dictionary
 
-        Returns:
-            Dictionary of new variables.
+        Args:
+            tensor_shape: tensor shape of the quantized tensor.
+            name: Tensor name.
+            layer: Layer to quantize.
         """
 
         # Add min and max variables to layer.
@@ -87,9 +82,9 @@ class STEUniformWeightQATQuantizer(BasePytorchQATTrainableQuantizer):
         layer.register_parameter(name+"_"+FQ_MAX, nn.Parameter(to_torch_tensor(self.max_values), requires_grad=False))
 
         # Save the quantizer parameters for later calculations
-        self.quantizer_parameters = {FQ_MIN: layer.get_parameter(name+"_"+FQ_MIN), FQ_MAX: layer.get_parameter(name+"_"+FQ_MAX)}
+        self.add_quantizer_variable(FQ_MIN, layer.get_parameter(name+"_"+FQ_MIN), VariableGroup.QPARAMS)
+        self.add_quantizer_variable(FQ_MAX, layer.get_parameter(name+"_"+FQ_MAX), VariableGroup.QPARAMS)
 
-        return self.quantizer_parameters
 
     def __call__(self,
                  inputs: nn.Parameter,
@@ -102,7 +97,7 @@ class STEUniformWeightQATQuantizer(BasePytorchQATTrainableQuantizer):
         Returns:
             quantized tensor
         """
-        return uniform_quantizer(inputs, self.quantizer_parameters[FQ_MIN], self.quantizer_parameters[FQ_MAX], self.num_bits)
+        return uniform_quantizer(inputs, self.get_quantizer_variable(FQ_MIN), self.get_quantizer_variable(FQ_MAX), self.num_bits)
 
     def convert2inferable(self) -> WeightsUniformInferableQuantizer:
         """
@@ -111,8 +106,8 @@ class STEUniformWeightQATQuantizer(BasePytorchQATTrainableQuantizer):
         Returns:
             A pytorch inferable quanizer object.
         """
-        _min = self.quantizer_parameters[FQ_MIN].cpu().detach().numpy()
-        _max = self.quantizer_parameters[FQ_MAX].cpu().detach().numpy()
+        _min = self.get_quantizer_variable(FQ_MIN).cpu().detach().numpy()
+        _max = self.get_quantizer_variable(FQ_MAX).cpu().detach().numpy()
 
         return WeightsUniformInferableQuantizer(num_bits=self.num_bits,
                                                 min_range=_min, max_range=_max,
@@ -143,21 +138,25 @@ class STEUniformActivationQATQuantizer(BasePytorchQATTrainableQuantizer):
         self.min_range_tensor = torch.Tensor([np_min_range])
         self.max_range_tensor = torch.Tensor([np_max_range])
         self.num_bits = quantization_config.activation_n_bits
-        self.quantizer_parameters = {}
 
     def initialize_quantization(self,
                                 tensor_shape: torch.Size,
                                 name: str,
-                                layer: qi.PytorchQuantizationWrapper) -> Dict[str, nn.Parameter]:
+                                layer: qi.PytorchQuantizationWrapper):
         """
-        Add min and max variables to layer.
+        Add quantizer parameters to the quantizer parameters dictionary
+
+        Args:
+            tensor_shape: tensor shape of the quantized tensor.
+            name: Tensor name.
+            layer: Layer to quantize.
         """
         layer.register_parameter(name+"_"+FQ_MIN, nn.Parameter(to_torch_tensor(self.min_range_tensor), requires_grad=True))
         layer.register_parameter(name+"_"+FQ_MAX, nn.Parameter(to_torch_tensor(self.max_range_tensor), requires_grad=True))
 
         # Save the quantizer parameters for later calculations
-        self.quantizer_parameters = {FQ_MIN: layer.get_parameter(name+"_"+FQ_MIN), FQ_MAX: layer.get_parameter(name+"_"+FQ_MAX)}
-        return self.quantizer_parameters
+        self.add_quantizer_variable(FQ_MIN, layer.get_parameter(name+"_"+FQ_MIN), VariableGroup.QPARAMS)
+        self.add_quantizer_variable(FQ_MAX, layer.get_parameter(name+"_"+FQ_MAX), VariableGroup.QPARAMS)
 
     def __call__(self,
                  inputs: torch.Tensor,
@@ -172,8 +171,8 @@ class STEUniformActivationQATQuantizer(BasePytorchQATTrainableQuantizer):
             The quantized tensor.
         """
 
-        _min = self.quantizer_parameters[FQ_MIN]
-        _max = self.quantizer_parameters[FQ_MAX]
+        _min = self.get_quantizer_variable(FQ_MIN)
+        _max = self.get_quantizer_variable(FQ_MAX)
         q_tensor = uniform_quantizer(inputs, _min, _max, self.num_bits)
         return q_tensor
 
@@ -184,8 +183,8 @@ class STEUniformActivationQATQuantizer(BasePytorchQATTrainableQuantizer):
         Returns:
             A pytorch inferable quanizer object.
         """
-        _min = self.quantizer_parameters[FQ_MIN].cpu().detach().numpy()
-        _max = self.quantizer_parameters[FQ_MAX].cpu().detach().numpy()
+        _min = self.get_quantizer_variable(FQ_MIN).cpu().detach().numpy()
+        _max = self.get_quantizer_variable(FQ_MAX).cpu().detach().numpy()
 
         return ActivationUniformInferableQuantizer(num_bits=self.num_bits,
                                                    min_range=_min, max_range=_max)

model_compression_toolkit/quantizers_infrastructure/inferable_infrastructure/common/get_all_subclasses.py

@@ -28,5 +28,4 @@ def get_all_subclasses(cls: type) -> Set[type]:
 
     """
 
-    return set(cls.__subclasses__()).union(
-        [s for c in cls.__subclasses__() for s in get_all_subclasses(c)])
+    return set(cls.__subclasses__()).union([s for c in cls.__subclasses__() for s in get_all_subclasses(c)])

model_compression_toolkit/quantizers_infrastructure/inferable_infrastructure/common/get_quantizers.py

@@ -41,7 +41,7 @@ def get_inferable_quantizer_class(quant_target: QuantizationTarget,
     qat_quantizer_classes = get_all_subclasses(quantizer_base_class)
     filtered_quantizers = list(filter(lambda q_class: getattr(q_class, QUANTIZATION_TARGET) == quant_target and
                                                       getattr(q_class, QUANTIZATION_METHOD) is not None and
-                                                       quant_method in getattr(q_class, QUANTIZATION_METHOD),
+                                                      quant_method in getattr(q_class, QUANTIZATION_METHOD),
                                       qat_quantizer_classes))
 
     if len(filtered_quantizers) != 1:

model_compression_toolkit/quantizers_infrastructure/inferable_infrastructure/keras/quantizers/__init__.py

@@ -18,6 +18,10 @@ from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructur
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers.weights_inferable_quantizers.weights_pot_inferable_quantizer import WeightsPOTInferableQuantizer
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers.weights_inferable_quantizers.weights_symmetric_inferable_quantizer import WeightsSymmetricInferableQuantizer
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers.weights_inferable_quantizers.weights_uniform_inferable_quantizer import WeightsUniformInferableQuantizer
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers.weights_inferable_quantizers.weights_lut_symmetric_inferable_quantizer import WeightsLUTSymmetricInferableQuantizer
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers.weights_inferable_quantizers.weights_lut_pot_inferable_quantizer import WeightsLUTPOTInferableQuantizer
+
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers.activation_inferable_quantizers.activation_pot_inferable_quantizer import ActivationPOTInferableQuantizer
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers.activation_inferable_quantizers.activation_symmetric_inferable_quantizer import ActivationSymmetricInferableQuantizer
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers.activation_inferable_quantizers.activation_uniform_inferable_quantizer import ActivationUniformInferableQuantizer
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers.activation_inferable_quantizers.activation_lut_pot_inferable_quantizer import ActivationLutPOTInferableQuantizer

model_compression_toolkit/quantizers_infrastructure/inferable_infrastructure/keras/quantizers/constants.py

@@ -22,3 +22,4 @@ MIN_RANGE = 'min_range'
 MAX_RANGE = 'max_range'
 CHANNEL_AXIS = 'channel_axis'
 INPUT_RANK = 'input_rank'
+CLUSTER_CENTERS = 'cluster_centers'

model_compression_toolkit/quantizers_infrastructure/inferable_infrastructure/pytorch/quantize_wrapper.py

@@ -184,13 +184,11 @@ if FOUND_TORCH:
             return self._weights_vars
 
         def forward(self,
-                    x: torch.Tensor,
                     *args: List[Any],
                     **kwargs: Dict[str, Any]) -> Union[torch.Tensor, List[torch.Tensor]]:
             """
             PytorchQuantizationWrapper forward functions
             Args:
-                x: layer's inputs
                 args: arguments to pass to internal layer.
                 kwargs: key-word dictionary to pass to the internal layer.
 
@@ -218,7 +216,7 @@ if FOUND_TORCH:
             # ----------------------------------
             # Layer operation
             # ----------------------------------
-            outputs = self.layer(x, *args, **kwargs)
+            outputs = self.layer(*args, **kwargs)
 
             # ----------------------------------
             # Quantize all activations
@@ -240,6 +238,18 @@ if FOUND_TORCH:
 
             return outputs
 
+        def get_quantized_weights(self) -> Dict[str, torch.Tensor]:
+            """
+
+            Returns: A dictionary of weights attributes to quantized weights.
+
+            """
+            quantized_weights = {}
+            weights_var = self.get_weights_vars()
+            for name, w, quantizer in weights_var:
+                quantized_weights[name] = quantizer(w)
+            return quantized_weights
+
 else:
     class PytorchQuantizationWrapper(object):
         def __init__(self,

model_compression_toolkit/quantizers_infrastructure/inferable_infrastructure/pytorch/quantizers/__init__.py

@@ -19,6 +19,8 @@ from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructur
     import ActivationSymmetricInferableQuantizer
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.pytorch.quantizers.activation_inferable_quantizers.activation_uniform_inferable_quantizer \
     import ActivationUniformInferableQuantizer
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.pytorch.quantizers.activation_inferable_quantizers.activation_lut_pot_inferable_quantizer \
+    import ActivationLutPOTInferableQuantizer
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.pytorch.quantizers.base_pytorch_inferable_quantizer \
     import BasePyTorchInferableQuantizer
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.pytorch.quantizers.weights_inferable_quantizers.weights_pot_inferable_quantizer \
@@ -27,3 +29,7 @@ from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructur
     import WeightsSymmetricInferableQuantizer
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.pytorch.quantizers.weights_inferable_quantizers.weights_uniform_inferable_quantizer \
     import WeightsUniformInferableQuantizer
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.pytorch.quantizers.weights_inferable_quantizers.weights_lut_symmetric_inferable_quantizer \
+    import WeightsLUTSymmetricInferableQuantizer
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.pytorch.quantizers.weights_inferable_quantizers.weights_lut_pot_inferable_quantizer \
+    import WeightsLUTPOTInferableQuantizer

model_compression_toolkit/quantizers_infrastructure/inferable_infrastructure/pytorch/quantizers/constants.py

@@ -21,3 +21,6 @@ PER_CHANNEL = 'per_channel'
 MIN_RANGE = 'min_range'
 MAX_RANGE = 'max_range'
 CHANNEL_AXIS = 'channel_axis'
+CLUSTER_CENTERS = 'cluster_centers'
+MULTIPLIER_N_BITS = 'multiplier_n_bits'
+EPS = 'eps'

model_compression_toolkit/quantizers_infrastructure/trainable_infrastructure/common/base_trainable_quantizer.py

@@ -12,8 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
-from typing import Union
+from abc import abstractmethod
+from enum import Enum
+from typing import Union, List, Any
 from inspect import signature
 
 from model_compression_toolkit.core import common
@@ -27,6 +28,19 @@ from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructur
     QUANTIZATION_TARGET
 
 
+VAR = 'var'
+GROUP = 'group'
+
+class VariableGroup(Enum):
+    """
+    An enum for choosing trainable variable group
+    0. WEIGHTS
+    1. QPARAMS
+    """
+    WEIGHTS = 0
+    QPARAMS = 1
+
+
 class BaseTrainableQuantizer(BaseInferableQuantizer):
     def __init__(self,
                  quantization_config: Union[TrainableQuantizerActivationConfig, TrainableQuantizerWeightsConfig]):
@@ -70,6 +84,8 @@ class BaseTrainableQuantizer(BaseInferableQuantizer):
             common.Logger.error(
                 f'Unknown Quantization Part:{static_quantization_target}')  # pragma: no cover
 
+        self.quantizer_parameters = {}
+
     @classmethod
     def get_sig(cls):
         return signature(cls)
@@ -147,3 +163,38 @@ class BaseTrainableQuantizer(BaseInferableQuantizer):
             BaseInferableQuantizer object.
         """
         raise NotImplemented  # pragma: no cover
+
+    def add_quantizer_variable(self, name: str, variable: Any, group: VariableGroup = VariableGroup.WEIGHTS):
+        """
+        Add a quantizer variable to quantizer_parameters dictionary
+        """
+        self.quantizer_parameters.update({name: {VAR: variable, GROUP: group}})
+
+    def get_quantizer_variable(self, name: str) -> Any:
+        """
+        Get a quantizer variable by name
+
+        Args:
+            name: variable name
+
+        Returns:
+            trainable variable
+        """
+        if name in self.quantizer_parameters:
+            return self.quantizer_parameters[name][VAR]
+        else:
+            common.Logger.error(f'Variable {name} is not exist in quantizers parameters!') # pragma: no cover
+
+
+    @abstractmethod
+    def get_trainable_variables(self, group: VariableGroup) -> List[Any]:
+        """
+        Get trainable parameters with specific group from quantizer
+
+        Args:
+            group: Enum of variable group
+
+        Returns:
+            List of trainable variables
+        """
+        raise NotImplemented  # pragma: no cover

model_compression_toolkit/quantizers_infrastructure/trainable_infrastructure/common/get_quantizers.py

@@ -14,7 +14,8 @@
 # ==============================================================================
 from typing import Union
 
-from model_compression_toolkit import TrainingMethod, RoundingType
+from model_compression_toolkit.gptq import RoundingType
+from model_compression_toolkit import TrainingMethod
 from model_compression_toolkit.core.common import Logger
 from model_compression_toolkit.core.common.target_platform import QuantizationMethod
 from model_compression_toolkit.quantizers_infrastructure import QuantizationTarget

model_compression_toolkit/quantizers_infrastructure/trainable_infrastructure/keras/base_keras_quantizer.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Dict, Any, Union
+from typing import Dict, Any, Union, List
 
 from model_compression_toolkit.core.common import Logger
 from model_compression_toolkit.core.common.constants import FOUND_TF
-
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import BaseTrainableQuantizer
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import BaseTrainableQuantizer, VAR, GROUP
 from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig, \
     TrainableQuantizerActivationConfig
 
@@ -25,7 +25,7 @@ if FOUND_TF:
     QUANTIZATION_CONFIG = 'quantization_config'
     from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.keras.config_serialization import config_serialization, \
         config_deserialization
-
+    import tensorflow as tf
 
     class BaseKerasTrainableQuantizer(BaseTrainableQuantizer):
         def __init__(self,
@@ -61,6 +61,24 @@ if FOUND_TF:
             # Note that a quantizer only receive quantization config and the rest of define hardcoded inside the speficie quantizer.
             return cls(quantization_config=quantization_config)
 
+        def get_trainable_variables(self, group: VariableGroup) -> List[tf.Tensor]:
+            """
+            Get trainable parameters with specific group from quantizer
+
+            Args:
+                group: Enum of variable group
+
+            Returns:
+                List of trainable variables
+            """
+            quantizer_trainable = []
+            for name, parameter_dict in self.quantizer_parameters.items():
+                quantizer_parameter, parameter_group = parameter_dict[VAR], parameter_dict[GROUP]
+                if quantizer_parameter.trainable and parameter_group == group:
+                    quantizer_trainable.append(quantizer_parameter)
+            return quantizer_trainable
+
+
 else:
     class BaseKerasTrainableQuantizer(BaseTrainableQuantizer):
         def __init__(self,

model_compression_toolkit/quantizers_infrastructure/trainable_infrastructure/pytorch/base_pytorch_quantizer.py

@@ -12,17 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Union
+from typing import Union, List
 
 from model_compression_toolkit.core.common.logger import Logger
 from model_compression_toolkit.core.common.constants import FOUND_TORCH
-
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import BaseTrainableQuantizer
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import BaseTrainableQuantizer, VAR, GROUP
 from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig, \
     TrainableQuantizerActivationConfig
 
+
 if FOUND_TORCH:
 
+    import torch
+
     class BasePytorchTrainableQuantizer(BaseTrainableQuantizer):
         def __init__(self,
                      quantization_config: Union[TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig]):
@@ -35,6 +38,24 @@ if FOUND_TORCH:
             """
             super().__init__(quantization_config)
 
+
+        def get_trainable_variables(self, group: VariableGroup) -> List[torch.Tensor]:
+            """
+            Get trainable parameters with specific group from quantizer
+
+            Args:
+                group: Enum of variable group
+
+            Returns:
+                List of trainable variables
+            """
+            quantizer_trainable = []
+            for name, parameter_dict in self.quantizer_parameters.items():
+                quantizer_parameter, parameter_group = parameter_dict[VAR], parameter_dict[GROUP]
+                if quantizer_parameter.requires_grad and parameter_group == group:
+                    quantizer_trainable.append(quantizer_parameter)
+            return quantizer_trainable
+
 else:
     class BasePytorchTrainableQuantizer(BaseTrainableQuantizer):
         def __init__(self,

model_compression_toolkit/gptq/common/gptq_quantizer_config.py

@@ -1,93 +0,0 @@
-# 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 Any, List, Callable
-
-from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT, REGULARIZATION_VALUES
-
-
-class GPTQQuantizerConfig:
-    """
-    A base class to define specific quantizer configuration for GPTQ quantizer.
-    """
-
-    def __init__(self):
-        self.n_batches = None
-
-    def get_regularization_value(self, fxp_model: Any, **kwargs) -> float:
-        """
-        Computes a regularization value for the quantizer's loss (if needed).
-        In the base class it only returns 0, to be used for GPTQ quantizers that don't require regularization.
-
-        Args:
-            fxp_model: The quantized model that is being trained.
-            **kwargs: Additional arguments for the quantizer regularization computation.
-
-        Returns: The regularization value.
-        """
-
-        return 0
-
-    def set_num_batches(self, num_batches: int):
-        """
-        Allows to set the number of batches that the quantizer uses for training (in each epoch).
-
-        Args:
-            num_batches: number of batches to be set.
-
-        """
-        self.n_batches = num_batches
-
-
-class SoftQuantizerConfig(GPTQQuantizerConfig):
-    def __init__(self, entropy_regularization: float = REG_DEFAULT):
-        """
-        Initializes an object that holds the arguments that are needed for soft rounding quantizer.
-
-        Args:
-            entropy_regularization (float): A floating point number that defines the gumbel entropy regularization factor.
-        """
-
-        super().__init__()
-        self.entropy_regularization = entropy_regularization
-
-
-    def get_regularization_value(self, fxp_model: Any, **kwargs) -> float:
-        """
-        Computes a regularization value for the soft quantizer.
-
-        Args:
-            fxp_model: The quantized model that is being trained.
-            **kwargs: Additional arguments for the quantizer regularization computation.
-
-        Returns: The regularization value.
-        """
-
-        soft_rounding_reg_values = kwargs.get(REGULARIZATION_VALUES)
-
-        if soft_rounding_reg_values is None:
-            Logger.error("No regularization values has been provided for computing the regularization "  # pragma: no cover
-                         "of the soft quantizer.")
-        if not isinstance(soft_rounding_reg_values, List):
-            Logger.error("The provided regularization values parameter of the soft quantizer "  # pragma: no cover
-                         "is not compatible (should be a list).")
-
-        reg = 0
-
-        for sq in soft_rounding_reg_values:
-            reg += sq
-
-        return self.entropy_regularization * reg
-