mct-nightly 1.8.0.22032023.post333__py3-none-any.whl → 1.8.0.22052023.post408__py3-none-any.whl

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
@@ -188,6 +188,13 @@ class SymmetricSoftRoundingGPTQ(BaseKerasGPTQTrainableQuantizer):
 
         ptq_threshold_tensor = self.get_quantizer_variable(PTQ_THRESHOLD)
 
+        #####################################################
+        # 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,
@@ -197,13 +204,6 @@ class SymmetricSoftRoundingGPTQ(BaseKerasGPTQTrainableQuantizer):
             # Calculate soft rounding targets and optimized threshold
             ##########################################################
             ptq_threshold_tensor_hat = tf.reshape(ptq_threshold_tensor, reshape_shape)
-            aux_var = self.get_soft_targets()
-
-            #####################################################
-            # Soft Rounding
-            #####################################################
-            if not training:
-                aux_var = tf.cast(tf.math.greater_equal(aux_var, 0.5), tf.float32)
 
             #####################################################
             # Quantized Input
@@ -219,14 +219,19 @@ class SymmetricSoftRoundingGPTQ(BaseKerasGPTQTrainableQuantizer):
                 scale = tf.reshape(self.get_quantizer_variable(SCALE_PTQ), reshape_shape)
                 q_tensor *= scale
 
-            return q_tensor
         else:
-            return soft_rounding_symmetric_quantizer(input_tensor=inputs,
-                                                     auxvar_tensor=self.quantizer_parameters[AUXVAR]['var'],
-                                                     threshold_tensor=ptq_threshold_tensor.value(),
-                                                     num_bits=self.num_bits,
-                                                     signed=True,
-                                                     power_of_two=self.power_of_two)
+            q_tensor = soft_rounding_symmetric_quantizer(input_tensor=inputs,
+                                                         auxvar_tensor=aux_var,
+                                                         threshold_tensor=ptq_threshold_tensor.value(),
+                                                         num_bits=self.num_bits,
+                                                         signed=True,
+                                                         power_of_two=self.power_of_two)
+
+            if self.quantization_parameter_learning and not self.power_of_two:
+                scale = self.get_quantizer_variable(SCALE_PTQ)
+                q_tensor *= scale
+
+        return q_tensor
 
     def get_quant_config(self) -> Dict[str, np.ndarray]:
         """

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/__init__.py

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

model_compression_toolkit/gptq/pytorch/quantizer/base_pytorch_gptq_quantizer.py

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

model_compression_toolkit/gptq/pytorch/quantizer/quant_utils.py

@@ -14,9 +14,7 @@
 # ==============================================================================
 from typing import Union, Tuple
 import torch
-from torch.nn.functional import softmax, log_softmax, one_hot
-from model_compression_toolkit.core.common.constants import MIN_THRESHOLD
-from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
+from model_compression_toolkit.constants import MIN_THRESHOLD
 
 
 def power_of_two_max(max_tensor: torch.Tensor) -> torch.Tensor:
@@ -30,11 +28,20 @@ def calculate_delta(max_tensor: torch.Tensor,
                     num_bits: int,
                     signed: bool) -> torch.Tensor:
     """
-    Compute the step size for the quantization.
+    Compute the step size for the symmetric quantization.
     """
     return max_tensor / (2 ** (num_bits - int(signed)))
 
 
+def calculate_delta_uniform(min_tensor: torch.Tensor,
+                            max_tensor: torch.Tensor,
+                            num_bits: int) -> torch.Tensor:
+    """
+    Compute the step size for the uniform quantization.
+    """
+    return (max_tensor-min_tensor) / (2 ** num_bits - 1)
+
+
 def ste_ceil(x: torch.Tensor) -> torch.Tensor:
     """
     Return the ceil values of a tensor.
@@ -42,6 +49,13 @@ def ste_ceil(x: torch.Tensor) -> torch.Tensor:
     return (torch.ceil(x) - x).detach() + x
 
 
+def ste_floor(x: torch.Tensor) -> torch.Tensor:
+    """
+    Return the floor values of a tensor.
+    """
+    return (torch.floor(x) - x).detach() + x
+
+
 def ste_round(x: torch.Tensor) -> torch.Tensor:
     """
     Calculate the rounded values of a tensor

model_compression_toolkit/gptq/pytorch/quantizer/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

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

@@ -19,7 +19,7 @@ import numpy as np
 
 from model_compression_toolkit.core.common import max_power_of_two
 from model_compression_toolkit import quantizers_infrastructure as qi
-from model_compression_toolkit.core.common.target_platform import QuantizationMethod
+from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
 from model_compression_toolkit.gptq.common.gptq_config import RoundingType
 from model_compression_toolkit.gptq.pytorch.quantizer.base_pytorch_gptq_quantizer import \
     BasePytorchGPTQTrainableQuantizer
@@ -27,7 +27,7 @@ from model_compression_toolkit.core.pytorch.utils import to_torch_tensor, torch_
 from model_compression_toolkit.gptq.pytorch.quantizer import quant_utils as qutils
 from model_compression_toolkit.gptq.common.gptq_constants import PTQ_THRESHOLD, SCALE_PTQ, \
     SOFT_ROUNDING_GAMMA, SOFT_ROUNDING_ZETA, AUXVAR
-from model_compression_toolkit.core.common.constants import THRESHOLD, MIN_THRESHOLD
+from model_compression_toolkit.constants import THRESHOLD, MIN_THRESHOLD
 from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.base_inferable_quantizer import mark_quantizer
 from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.quant_utils import \
@@ -142,9 +142,13 @@ class SymmetricSoftRoundingGPTQ(BasePytorchGPTQTrainableQuantizer):
         self.add_quantizer_variable(AUXVAR, layer.get_parameter(f"{name}_{AUXVAR}"), VariableGroup.WEIGHTS)
 
         if self.quantization_parameter_learning:
-            layer.register_parameter(f"{name}_{SCALE_PTQ}",
-                                     nn.Parameter(to_torch_tensor(torch.ones_like(torch.Tensor(self.threshold_values))),
-                                                  requires_grad=True))
+            if self.per_channel:
+                layer.register_parameter(f"{name}_{SCALE_PTQ}",
+                                         nn.Parameter(to_torch_tensor(torch.ones_like(torch.Tensor(self.threshold_values))),
+                                                      requires_grad=True))
+            else:
+                layer.register_parameter(f"{name}_{SCALE_PTQ}",
+                                         nn.Parameter(to_torch_tensor((torch.tensor([1.0], requires_grad=True)))))
             self.add_quantizer_variable(SCALE_PTQ, layer.get_parameter(f"{name}_{SCALE_PTQ}"), VariableGroup.QPARAMS)
 
     def get_soft_targets(self) -> torch.Tensor:
@@ -233,4 +237,8 @@ class SymmetricSoftRoundingGPTQ(BasePytorchGPTQTrainableQuantizer):
                                                          signed=True,
                                                          power_of_two=self.power_of_two)
 
+            if self.quantization_parameter_learning and not self.power_of_two:
+                scale = self.get_quantizer_variable(SCALE_PTQ)
+                q_tensor *= scale
+
         return q_tensor