mct-nightly 1.8.0.8042023.post345__py3-none-any.whl → 1.8.0.8052023.post414__py3-none-any.whl

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.
@@ -94,9 +94,16 @@ if common.constants.FOUND_TF:
             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)
-        return GradientPTQConfigV2(n_epochs, optimizer, optimizer_rest=optimizer_rest, loss=loss,
-                                   log_function=log_function, train_bias=True, optimizer_bias=bias_optimizer)
+        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,
+                                   optimizer_bias=bias_optimizer,
+                                   use_hessian_based_weights=use_hessian_based_weights)
 
 
     def keras_gradient_post_training_quantization_experimental(in_model: Model,
@@ -158,20 +165,20 @@ 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:
 
@@ -187,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, \

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/soft_rounding/soft_quantizer_reg.py

@@ -77,7 +77,7 @@ class SoftQuantizerRegularization:
         # Initializing the temperature decay according to the number of expected gradient steps
         self.linear_decay = LinearTempDecay(total_gradient_steps)
 
-        self.count_iter = 0
+        self.count_iter = tf.Variable(0.)
 
 
     def __call__(self, model: Model, entropy_reg: float):
@@ -90,16 +90,14 @@ class SoftQuantizerRegularization:
 
         Returns: Regularization value.
         """
-
         soft_reg_aux: List[tf.Tensor] = []
+        b = self.linear_decay(self.count_iter.value())
         for layer in model.layers:
             if isinstance(layer, KerasQuantizationWrapper):
                 kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=type(layer.layer),
                                                                       fw_info=DEFAULT_KERAS_INFO)
 
                 st = layer.weights_quantizers[kernel_attribute].get_soft_targets()
-                b = self.linear_decay(self.count_iter)
-
                 soft_reg_aux.append(tf.reduce_sum(1 - tf.pow(tf.math.abs(st - .5) * 2, b)))
 
         reg = 0
@@ -107,6 +105,6 @@ class SoftQuantizerRegularization:
         for sq in soft_reg_aux:
             reg += sq
 
-        self.count_iter += 1
+        self.count_iter.assign_add(1.0)
 
         return entropy_reg * reg

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

@@ -19,12 +19,12 @@ import numpy as np
 from model_compression_toolkit.gptq import RoundingType
 from model_compression_toolkit import quantizers_infrastructure as qi
 from model_compression_toolkit.core.common import max_power_of_two
-from model_compression_toolkit.core.common.target_platform import QuantizationMethod
+from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
 from model_compression_toolkit.gptq.common.gptq_constants import PTQ_THRESHOLD, SCALE_PTQ, \
     SOFT_ROUNDING_GAMMA, SOFT_ROUNDING_ZETA, AUXVAR
 from model_compression_toolkit.gptq.keras.quantizer import quant_utils as qutils
 from typing import Dict, Any
-from model_compression_toolkit.core.common.constants import THRESHOLD, MIN_THRESHOLD
+from model_compression_toolkit.constants import THRESHOLD, MIN_THRESHOLD
 from model_compression_toolkit.gptq.keras.quantizer.base_keras_gptq_quantizer import BaseKerasGPTQTrainableQuantizer
 from model_compression_toolkit.gptq.keras.quantizer.quant_utils import power_of_two_max, clip, calculate_delta
 from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig

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

@@ -0,0 +1,224 @@
+# Copyright 2023 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import tensorflow as tf
+import numpy as np
+
+from model_compression_toolkit.gptq import RoundingType
+from model_compression_toolkit import quantizers_infrastructure as qi
+from model_compression_toolkit.quantizers_infrastructure.constants import FQ_MIN, FQ_MAX
+from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
+from model_compression_toolkit.gptq.common.gptq_constants import \
+    SOFT_ROUNDING_GAMMA, SOFT_ROUNDING_ZETA, AUXVAR
+from model_compression_toolkit.gptq.keras.quantizer import quant_utils as qutils
+from typing import Dict, Any
+from model_compression_toolkit.constants import RANGE_MIN, RANGE_MAX
+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 soft_rounding_uniform_quantizer(input_tensor: tf.Tensor,
+                                    auxvar_tensor: tf.Variable,
+                                    min_tensor: tf.Tensor,
+                                    max_tensor: tf.Tensor,
+                                    num_bits: int) -> tf.Tensor:
+    """
+    Quantize a tensor uniformly for GPTQ quantizers.
+
+    Args:
+        input_tensor: Tensor to quantize. values of this tensor are not changed during gptq.
+        auxvar_tensor: Tensor that manifests the bit shift of the quantized weights due to gptq training.
+        min_tensor: Tensor with values to compute the min threshold.
+        max_tensor: Tensor with values to compute the max threshold.
+        num_bits: Num of bits to use.
+
+    Returns:
+        A quantized tensor.
+    """
+    # adjusts the quantization range so the quantization grid includes zero.
+    min_range, max_range = qutils.fix_range_to_include_zero(min_tensor, max_tensor, num_bits)
+    delta = qutils.calculate_delta_uniform(min_range, max_range, num_bits)
+    input_tensor_int = qutils.ste_floor((input_tensor - min_range) / delta)
+    tensor_q = input_tensor_int + auxvar_tensor
+    return delta * qutils.ste_clip(tensor_q,
+                                   min_val=0,
+                                   max_val=2 ** num_bits - 1) + min_range
+
+
+@mark_quantizer(quantization_target=qi.QuantizationTarget.Weights,
+                quantization_method=[QuantizationMethod.UNIFORM],
+                quantizer_type=RoundingType.SoftQuantizer)
+class UniformSoftRoundingGPTQ(BaseKerasGPTQTrainableQuantizer):
+    """
+    Trainable uniform quantizer to optimize the rounding of the quantized values using a soft quantization method.
+    """
+
+    def __init__(self,
+                 quantization_config: TrainableQuantizerWeightsConfig,
+                 quantization_parameter_learning: bool = False):
+        """
+        Initialize a UniformSoftRoundingGPTQ object with parameters to use
+        for the quantization.
+
+        Args:
+            quantization_config: Trainable weight quantizer config.
+            quantization_parameter_learning: Whether to train the quantization threshold.
+        """
+        super().__init__(quantization_config)
+        self.num_bits = quantization_config.weights_n_bits
+        self.per_channel = quantization_config.weights_per_channel_threshold
+
+        self.min_values = quantization_config.weights_quantization_params[RANGE_MIN]
+        self.max_values = quantization_config.weights_quantization_params[RANGE_MAX]
+
+        self.quantization_axis = quantization_config.weights_channels_axis
+        assert quantization_parameter_learning is False, \
+            "Quantization parameters learning in UniformSoftRoundingGPTQ not implemented yet"
+        self.quantization_parameter_learning = quantization_parameter_learning
+        self.num_channels = self.min_values.shape[self.quantization_axis] if self.per_channel else 1
+
+        # gamma and zeta are stretch parameters for computing the rectified sigmoid function.
+        # See: https://arxiv.org/pdf/2004.10568.pdf
+        self.gamma = SOFT_ROUNDING_GAMMA
+        self.zeta = SOFT_ROUNDING_ZETA
+
+    def initialize_quantization(self,
+                                tensor_shape: Any,
+                                name: str,
+                                layer: Any):
+        """
+        Add quantizer parameters to the quantizer parameters dictionary
+
+        Args:
+            tensor_shape: tensor shape of the quantized tensor.
+            name: Tensor name.
+            layer: Layer to quantize.
+        """
+
+        if self.per_channel:
+            reshape_shape = get_threshold_reshape_shape(tensor_shape,
+                                                        quant_axis=self.quantization_axis,
+                                                        quant_axis_dim=self.num_channels)
+        else:
+            reshape_shape = [self.num_channels]
+
+        min_tensor = layer.add_weight(
+            f"{name}_{FQ_MIN}",
+            shape=reshape_shape,
+            initializer=tf.keras.initializers.Constant(1.0),
+            trainable=False)
+        min_tensor.assign(self.min_values.reshape(reshape_shape))
+
+        max_tensor = layer.add_weight(
+            f"{name}_{FQ_MAX}",
+            shape=reshape_shape,
+            initializer=tf.keras.initializers.Constant(1.0),
+            trainable=False)
+        max_tensor.assign(self.max_values.reshape(reshape_shape))
+
+        w = getattr(layer.layer, name)
+        auxvar_tensor = layer.add_weight(
+            f"{name}_{AUXVAR}",
+            shape=list(w.shape),
+            initializer=tf.keras.initializers.Constant(0.0),
+            trainable=True)
+
+        w = layer.layer.depthwise_kernel if isinstance(layer.layer, (tf.keras.layers.DepthwiseConv2D,
+                                                                     tf.keras.layers.DepthwiseConv1D)) \
+            else layer.layer.kernel
+        delta = qutils.calculate_delta_uniform(min_tensor, max_tensor, self.num_bits)
+        w_clipped_normed = qutils.clip((w - min_tensor)/ delta, 0, 2 ** self.num_bits - 1)
+        rest = w_clipped_normed - tf.floor(w_clipped_normed)  # rest of rounding [0, 1)
+        alpha = -qutils.safe_log((self.zeta - self.gamma) / (rest - self.gamma) - 1, 1e-16)  # => sigmoid(alpha) = rest
+        auxvar_tensor.assign(alpha)
+
+        # Add quantization variables
+        self.add_quantizer_variable(AUXVAR, auxvar_tensor, VariableGroup.WEIGHTS)
+        self.add_quantizer_variable(RANGE_MIN, min_tensor, VariableGroup.QPARAMS)
+        self.add_quantizer_variable(RANGE_MAX, max_tensor, VariableGroup.QPARAMS)
+
+    def get_soft_targets(self) -> tf.Tensor:
+        """
+        Computes the rectified sigmoid function for the quantization target parameters.
+
+        Returns:
+            A tensor with the soft rounding targets values.
+
+        """
+        return qutils.clip(
+            tf.sigmoid(self.get_quantizer_variable(AUXVAR)) * (self.zeta - self.gamma) + self.gamma, 1, 0)
+
+    def __call__(self,
+                 inputs: tf.Tensor,
+                 training: bool):
+        """
+        Quantize a tensor.
+
+        Args:
+            inputs: Input tensor to quantize.
+            training: Whether the graph is in training mode.
+
+        Returns:
+            The quantized tensor.
+        """
+
+        min_tensor = self.get_quantizer_variable(RANGE_MIN)
+        max_tensor = self.get_quantizer_variable(RANGE_MAX)
+
+        #####################################################
+        # Soft Rounding
+        #####################################################
+        aux_var = self.get_soft_targets()
+        if not training:
+            aux_var = tf.cast(tf.math.greater_equal(aux_var, 0.5), tf.float32)
+
+        if self.per_channel:
+            reshape_shape = get_threshold_reshape_shape(inputs.shape,
+                                                        quant_axis=self.quantization_axis,
+                                                        quant_axis_dim=-1)
+
+            #####################################################
+            # Quantized Input
+            #####################################################
+            q_tensor = soft_rounding_uniform_quantizer(input_tensor=inputs,
+                                                       auxvar_tensor=aux_var,
+                                                       min_tensor=tf.reshape(min_tensor, reshape_shape),
+                                                       max_tensor=tf.reshape(max_tensor, reshape_shape),
+                                                       num_bits=self.num_bits)
+
+        else:
+            q_tensor = soft_rounding_uniform_quantizer(input_tensor=inputs,
+                                                       auxvar_tensor=aux_var,
+                                                       min_tensor=min_tensor,
+                                                       max_tensor=max_tensor,
+                                                       num_bits=self.num_bits)
+
+        return q_tensor
+
+    def get_quant_config(self) -> Dict[str, np.ndarray]:
+        """
+        Returns the config used to edit NodeQuantizationConfig after GPTQ retraining
+
+        Returns:
+            A dictionary of attributes the quantize_config retraining has changed during GPTQ retraining.
+            Keys must match NodeQuantizationConfig attributes
+        """
+
+        return {RANGE_MIN: self.get_quantizer_variable(RANGE_MIN).numpy(),
+                RANGE_MAX: self.get_quantizer_variable(RANGE_MAX).numpy()}

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

@@ -20,10 +20,10 @@ import tensorflow as tf
 
 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.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

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,7 +19,7 @@ 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

model_compression_toolkit/gptq/pytorch/quantization_facade.py

@@ -14,18 +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.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.core.common.target_platform import TargetPlatformCapabilities
+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
+from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
     MixedPrecisionQuantizationConfigV2
 
@@ -36,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
@@ -118,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:
@@ -142,7 +142,7 @@ 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
 
@@ -152,16 +152,16 @@ if FOUND_TORCH:
 
         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
@@ -192,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/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, \

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:
@@ -51,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/soft_rounding/soft_quantizer_reg.py

@@ -95,14 +95,13 @@ class SoftQuantizerRegularization:
         """
 
         soft_reg_aux: List[torch.Tensor] = []
+        b = self.linear_decay(self.count_iter)
         for layer in model.modules():
             if isinstance(layer, PytorchQuantizationWrapper):
                 kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=type(layer.layer),
                                                                       fw_info=DEFAULT_PYTORCH_INFO)
 
                 st = layer.weights_quantizers[kernel_attribute].get_soft_targets()
-                b = self.linear_decay(self.count_iter)
-
                 soft_reg_aux.append((1 - torch.pow(torch.abs(st - .5) * 2, b)).sum())
 
         reg = 0