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

model_compression_toolkit/gptq/common/gptq_graph.py

@@ -14,8 +14,8 @@
 # ==============================================================================
 from typing import Tuple, List
 
-from model_compression_toolkit import FrameworkInfo
-from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.core import FrameworkInfo
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
 

model_compression_toolkit/gptq/common/gptq_training.py

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

model_compression_toolkit/gptq/keras/gptq_keras_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.keras.keras_implementation import KerasImplementation
+from model_compression_toolkit.gptq.common.gptq_framework_implementation import GPTQFrameworkImplemantation
+from model_compression_toolkit.gptq.keras.gptq_training import KerasGPTQTrainer
+
+
+class GPTQKerasImplemantation(GPTQFrameworkImplemantation, KerasImplementation):
+
+    def get_gptq_trainer_obj(self) -> Type[KerasGPTQTrainer]:
+        """
+        Returns:  Keras object of GPTQTrainer
+        """
+        return KerasGPTQTrainer

model_compression_toolkit/gptq/keras/gptq_training.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from functools import partial
 from typing import Callable, List, Tuple, Union
 
 import tensorflow as tf
@@ -23,11 +22,11 @@ from tqdm import tqdm
 # As from Tensorflow 2.6, keras is a separate package and some classes should be imported differently.
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder
-from model_compression_toolkit.gptq.common.gptq_constants import REGULARIZATION_VALUES
 from packaging import version
 
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
 from model_compression_toolkit.gptq.keras.quantizer.quantization_builder import quantization_builder
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.quantizers_infrastructure import KerasQuantizationWrapper
 
 if version.parse(tf.__version__) < version.parse("2.6"):
@@ -37,15 +36,15 @@ else:
 
 from model_compression_toolkit.core import common
 from model_compression_toolkit.gptq.common.gptq_training import GPTQTrainer
-from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfigV2, RoundingType
+from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfigV2
 from model_compression_toolkit.core.common import Graph
-from model_compression_toolkit.gptq.keras.graph_info import get_weights_for_loss, \
-    get_soft_rounding_reg, get_gptq_trainable_parameters
+from model_compression_toolkit.gptq.keras.graph_info import get_weights_for_loss, get_gptq_trainable_parameters
+from model_compression_toolkit.gptq.keras.quantizer.regularization_factory import get_regularization
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 import numpy as np
 import copy
-from model_compression_toolkit.core.keras.constants import BIAS, USE_BIAS, KERNEL
+from model_compression_toolkit.core.keras.constants import BIAS, USE_BIAS
 from model_compression_toolkit import quantizers_infrastructure as qi
 
 
@@ -79,13 +78,12 @@ class KerasGPTQTrainer(GPTQTrainer):
                          graph_quant,
                          gptq_config,
                          fw_impl,
-                         fw_info,
-                         representative_data_gen)
+                         fw_info)
 
         self.loss_list = []
         self.input_scale = 1
 
-        trainable_weights, bias_weights, trainable_threshold, temperature_weights = get_gptq_trainable_parameters(
+        trainable_weights, bias_weights, trainable_threshold = get_gptq_trainable_parameters(
             self.fxp_model,
             fw_info,
             add_bias=gptq_config.train_bias)
@@ -108,11 +106,13 @@ class KerasGPTQTrainer(GPTQTrainer):
             [len(optimizer_params_tuple[1]) for optimizer_params_tuple in self.optimizer_with_param]) > 0
 
         if self.float_user_info.input_scale != self.gptq_user_info.input_scale:
-            common.Logger.error("Input scale mismatch between float and GPTQ networks")  # pragma: no cover
+            Logger.error("Input scale mismatch between float and GPTQ networks")  # pragma: no cover
         else:
             self.input_scale = self.gptq_user_info.input_scale
 
-        self.weights_for_average_loss = self.compute_jacobian_based_weights(representative_data_gen)
+        self.weights_for_average_loss = self.compute_hessian_based_weights(representative_data_gen)
+
+        self.reg_func = get_regularization(self.gptq_config, representative_data_gen)
 
     def _is_gptq_applicable(self,
                             node: common.BaseNode) -> bool:
@@ -127,7 +127,7 @@ class KerasGPTQTrainer(GPTQTrainer):
         """
 
         if node.is_weights_quantization_enabled() and not self.fw_info.is_kernel_op(node.type):
-            common.Logger.error(f"GPTQ Error: Quantizing node {node.name} of type {node.type} "
+            Logger.error(f"GPTQ Error: Quantizing node {node.name} of type {node.type} "
                                 f"without a kernel isn't supported")
         return node.is_weights_quantization_enabled()
 
@@ -195,9 +195,7 @@ class KerasGPTQTrainer(GPTQTrainer):
                                                self.compare_points_std,
                                                self.weights_for_average_loss)
 
-            reg_value = self.gptq_config.quantizer_config.get_regularization_value(
-                self.fxp_model,
-                **{REGULARIZATION_VALUES: self._get_quantizer_regularization_values(self.gptq_config.rounding_type)})
+            reg_value = self.reg_func(self.fxp_model, self.gptq_config.regularization_factor)
 
             loss_value += reg_value
 
@@ -283,7 +281,7 @@ class KerasGPTQTrainer(GPTQTrainer):
                     self.gptq_config.log_function(loss_value_step, grads[0], in_optimizer_with_param[0][-1],
                                                   self.compare_points)
                 self.loss_list.append(loss_value_step.numpy())
-                common.Logger.debug(f'last loss value: {self.loss_list[-1]}')
+                Logger.debug(f'last loss value: {self.loss_list[-1]}')
 
     def update_graph(self):
         """
@@ -300,7 +298,7 @@ class KerasGPTQTrainer(GPTQTrainer):
                 if len(node) == 0 and isinstance(layer.layer, TensorFlowOpLayer):
                     node = graph.find_node_by_name('_'.join(layer.layer.name.split('_')[3:]))
                 if len(node) != 1:
-                    common.Logger.error(f"Can't update GPTQ graph due to missing layer named: {layer.layer.name}")
+                    Logger.error(f"Can't update GPTQ graph due to missing layer named: {layer.layer.name}")
                 node = node[0]
                 kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=node.type,
                                                                       fw_info=self.fw_info)
@@ -319,18 +317,3 @@ class KerasGPTQTrainer(GPTQTrainer):
                         node.set_weights_by_keys(BIAS, new_bias)
 
         return graph
-
-    def _get_quantizer_regularization_values(self, rounding_type: RoundingType) -> List[tf.Tensor]:
-        """
-        Mapping between a rounding type to its matching regularization method.
-
-        Args:
-            rounding_type: GPTQ rounding type.
-
-        Returns: A regularization computation method.
-
-        """
-        if rounding_type == RoundingType.SoftQuantizer:
-            return get_soft_rounding_reg(self.fxp_model)
-        else:
-            return []

model_compression_toolkit/gptq/keras/graph_info.py

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

model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -16,21 +16,19 @@
 from typing import Callable, Tuple
 from packaging import version
 
-from model_compression_toolkit.core import common
-from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.core.common.constants import TENSORFLOW
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import TENSORFLOW, FOUND_TF
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfigV2
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
-from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
-    MixedPrecisionQuantizationConfigV2
-from model_compression_toolkit import CoreConfig
+from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
 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.core.common.target_platform.targetplatform2framework import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import TargetPlatformCapabilities
 
 LR_DEFAULT = 0.15
 LR_REST_DEFAULT = 1e-4
@@ -38,14 +36,14 @@ LR_BIAS_DEFAULT = 1e-4
 LR_QUANTIZATION_PARAM_DEFAULT = 1e-3
 GPTQ_MOMENTUM = 0.9
 
-if common.constants.FOUND_TF:
+if FOUND_TF:
     import tensorflow as tf
     from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
-    from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
+    from model_compression_toolkit.gptq.keras.gptq_keras_implementation import GPTQKerasImplemantation
     from model_compression_toolkit.core.keras.keras_model_validation import KerasModelValidation
     from tensorflow.keras.models import Model
     from model_compression_toolkit.gptq.keras.gptq_loss import GPTQMultipleTensorsLoss
-    from model_compression_toolkit.core.keras.constants import DEFAULT_TP_MODEL
+    from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from model_compression_toolkit.exporter.model_wrapper import get_exportable_keras_model
     from model_compression_toolkit import get_target_platform_capabilities
 
@@ -62,7 +60,8 @@ if common.constants.FOUND_TF:
                               optimizer: OptimizerV2 = tf.keras.optimizers.Adam(learning_rate=LR_DEFAULT),
                               optimizer_rest: OptimizerV2 = tf.keras.optimizers.Adam(learning_rate=LR_REST_DEFAULT),
                               loss: Callable = GPTQMultipleTensorsLoss(),
-                              log_function: Callable = None) -> GradientPTQConfigV2:
+                              log_function: Callable = None,
+                              use_hessian_based_weights: bool = True) -> GradientPTQConfigV2:
         """
         Create a GradientPTQConfigV2 instance for Keras models.
 
@@ -72,6 +71,7 @@ if common.constants.FOUND_TF:
             optimizer_rest (OptimizerV2): Keras optimizer to use for fine-tuning of the bias variable.
             loss (Callable): loss to use during fine-tuning. should accept 4 lists of tensors. 1st list of quantized tensors, the 2nd list is the float tensors, the 3rd is a list of quantized weights and the 4th is a list of float weights.
             log_function (Callable): Function to log information about the gptq process.
+            use_hessian_based_weights (bool): Whether to use Hessian-based weights for weighted average loss.
 
         returns:
             a GradientPTQConfigV2 object to use when fine-tuning the quantized model using gptq.
@@ -85,24 +85,25 @@ if common.constants.FOUND_TF:
 
             Create a GradientPTQConfigV2 to run for 5 epochs:
 
-            >>> gptq_conf = mct.get_keras_gptq_config(n_epochs=5)
+            >>> gptq_conf = mct.gptq.get_keras_gptq_config(n_epochs=5)
 
             Other Tensorflow optimizers can be passed:
 
-            >>> gptq_conf = mct.get_keras_gptq_config(n_epochs=3, optimizer=tf.keras.optimizers.Nadam())
+            >>> gptq_conf = mct.gptq.get_keras_gptq_config(n_epochs=3, optimizer=tf.keras.optimizers.Nadam())
 
             The configuration can be passed to :func:`~model_compression_toolkit.keras_post_training_quantization` in order to quantize a keras model using gptq.
 
         """
-        bias_optimizer = tf.keras.optimizers.SGD(learning_rate=LR_BIAS_DEFAULT, momentum=GPTQ_MOMENTUM)
+        bias_optimizer = tf.keras.optimizers.SGD(learning_rate=LR_BIAS_DEFAULT,
+                                                 momentum=GPTQ_MOMENTUM)
         return GradientPTQConfigV2(n_epochs,
                                    optimizer,
                                    optimizer_rest=optimizer_rest,
                                    loss=loss,
                                    log_function=log_function,
                                    train_bias=True,
-                                   quantization_parameters_learning=True,
-                                   optimizer_bias=bias_optimizer)
+                                   optimizer_bias=bias_optimizer,
+                                   use_hessian_based_weights=use_hessian_based_weights)
 
 
     def keras_gradient_post_training_quantization_experimental(in_model: Model,
@@ -164,28 +165,28 @@ if common.constants.FOUND_TF:
 
             Create an MCT core config, containing the quantization configuration:
 
-            >>> config = mct.CoreConfig()
+            >>> config = mct.core.CoreConfig()
 
             If mixed precision is desired, create an MCT core config with a mixed-precision configuration, to quantize a model
             with different bitwidths for different layers.
             The candidates bitwidth for quantization should be defined in the target platform model:
 
-            >>> config = mct.CoreConfig(mixed_precision_config=mct.MixedPrecisionQuantizationConfigV2(num_of_images=1))
+            >>> config = mct.core.CoreConfig(mixed_precision_config=mct.core.MixedPrecisionQuantizationConfigV2(num_of_images=1))
 
             For mixed-precision set a target KPI object:
             Create a KPI object to limit our returned model's size. Note that this value affects only coefficients
             that should be quantized (for example, the kernel of Conv2D in Keras will be affected by this value,
             while the bias will not):
 
-            >>> kpi = mct.KPI(model.count_params() * 0.75)  # About 0.75 of the model size when quantized with 8 bits.
+            >>> kpi = mct.core.KPI(model.count_params() * 0.75)  # About 0.75 of the model size when quantized with 8 bits.
 
             Create GPTQ config:
 
-            >>> gptq_config = mct.get_keras_gptq_config(n_epochs=1)
+            >>> gptq_config = mct.gptq.get_keras_gptq_config(n_epochs=1)
 
             Pass the model with the representative dataset generator to get a quantized model:
 
-            >>> quantized_model, quantization_info = mct.keras_gradient_post_training_quantization_experimental(model, repr_datagen, gptq_config, target_kpi=kpi, core_config=config)
+            >>> quantized_model, quantization_info = mct.gptq.keras_gradient_post_training_quantization_experimental(model, repr_datagen, gptq_config, target_kpi=kpi, core_config=config)
 
         """
         KerasModelValidation(model=in_model,
@@ -193,15 +194,15 @@ if common.constants.FOUND_TF:
 
         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(fw_info)
 
-        fw_impl = KerasImplementation()
+        fw_impl = GPTQKerasImplemantation()
 
         tg, bit_widths_config = core_runner(in_model=in_model,
                                             representative_data_gen=representative_data_gen,

model_compression_toolkit/gptq/keras/quantizer/__init__.py

@@ -15,3 +15,4 @@
 
 import model_compression_toolkit.gptq.keras.quantizer.ste_rounding.symmetric_ste
 import model_compression_toolkit.gptq.keras.quantizer.soft_rounding.symmetric_soft_quantizer
+import model_compression_toolkit.gptq.keras.quantizer.soft_rounding.uniform_soft_quantizer

model_compression_toolkit/gptq/keras/quantizer/base_keras_gptq_quantizer.py

@@ -15,8 +15,8 @@
 from abc import abstractmethod
 from typing import Union, Dict, List
 
-from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.core.common.constants import FOUND_TF
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import FOUND_TF
 from model_compression_toolkit.gptq.common.gptq_constants import WEIGHTS_QUANTIZATION_PARAMS
 
 from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig, \
@@ -45,7 +45,6 @@ if FOUND_TF:
 
             super().__init__(quantization_config)
 
-            self.quantizer_parameters = None
 
         def update_layer_quantization_params(self, layer: KerasQuantizationWrapper
                                              ) -> (Dict[str, tf.Tensor], Dict[str, Dict], Dict):

model_compression_toolkit/gptq/keras/quantizer/quant_utils.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 
 import tensorflow as tf
-from model_compression_toolkit.core.common.constants import MIN_THRESHOLD
+from model_compression_toolkit.constants import MIN_THRESHOLD
 from typing import Tuple
 
 
@@ -26,6 +26,14 @@ def ste_ceil(x: tf.Tensor) -> tf.Tensor:
     return error + x
 
 
+def ste_floor(x: tf.Tensor) -> tf.Tensor:
+    """
+    Return the floor values of a tensor.
+    """
+    error = tf.stop_gradient(tf.math.floor(x) - x)
+    return error + x
+
+
 def safe_log(x: tf.Tensor, eps: float) -> tf.Tensor:
     """
     Computes log function of x unless x is smaller than some small value, so the log function would not fail.
@@ -72,6 +80,15 @@ def calculate_delta(max_tensor: tf.Tensor,
     return max_tensor / (2 ** (num_bits - int(signed)))
 
 
+def calculate_delta_uniform(min_tensor: tf.Tensor,
+                            max_tensor: tf.Tensor,
+                            num_bits: int) -> tf.Tensor:
+    """
+    Compute the step size for the uniform quantization.
+    """
+    return (max_tensor-min_tensor) / (2 ** num_bits - 1)
+
+
 def ste_clip(x: [tf.Tensor, tf.Variable], max_val=1, min_val=None) -> tf.Tensor:
     """
     clip a variable between fixed values such that min_val<=output<=max_val

model_compression_toolkit/gptq/keras/quantizer/quantization_builder.py

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