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

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

@@ -0,0 +1,194 @@
+# 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 torch
+import torch.nn as nn
+from typing import Dict
+import numpy as np
+
+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_config import RoundingType
+from model_compression_toolkit.gptq.pytorch.quantizer.base_pytorch_gptq_quantizer import \
+    BasePytorchGPTQTrainableQuantizer
+from model_compression_toolkit.core.pytorch.utils import to_torch_tensor, torch_tensor_to_numpy
+from model_compression_toolkit.gptq.pytorch.quantizer import quant_utils as qutils
+from model_compression_toolkit.gptq.common.gptq_constants import SOFT_ROUNDING_GAMMA, SOFT_ROUNDING_ZETA, AUXVAR
+from model_compression_toolkit.gptq.pytorch.quantizer.quant_utils import fix_range_to_include_zero
+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.base_trainable_quantizer import \
+    VariableGroup
+from model_compression_toolkit.constants import RANGE_MAX, RANGE_MIN
+
+
+def soft_rounding_unifrom_quantizer(input_tensor: torch.Tensor,
+                                    auxvar_tensor: torch.Tensor,
+                                    min_range: torch.Tensor,
+                                    max_range: torch.Tensor,
+                                    num_bits: int) -> torch.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_range: Tensor with min values to compute the delta grid.
+        max_range: Tensor with max values to compute the delta grid.
+        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 = fix_range_to_include_zero(min_range, max_range, 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(BasePytorchGPTQTrainableQuantizer):
+    """
+    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):
+        """
+        Construct a Pytorch model that utilize a fake weight quantizer of soft-quantizer for symmetric quantizer.
+
+        Args:
+            quantization_config: Trainable weights quantizer config.
+            quantization_parameter_learning (Bool): Whether to learn the min/max ranges or not
+        """
+
+        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
+        self.quantization_parameter_learning = quantization_parameter_learning
+
+        # 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: torch.Size,
+                                name: str,
+                                layer: qi.PytorchQuantizationWrapper):
+        """
+        Add quantizer parameters to the quantizer parameters dictionary
+
+        Args:
+            tensor_shape: tensor shape of the quantized tensor.
+            name: Tensor name.
+            layer: Layer to quantize.
+        """
+
+        # Add min and max variables to layer.
+        if self.per_channel:
+            min_values = to_torch_tensor(self.min_values)
+            max_values = to_torch_tensor(self.max_values)
+        else:
+            min_values = torch.tensor(self.min_values)
+            max_values = torch.tensor(self.max_values)
+
+        layer.register_parameter(name+"_"+FQ_MIN, nn.Parameter(min_values, requires_grad=self.quantization_parameter_learning))
+        layer.register_parameter(name+"_"+FQ_MAX, nn.Parameter(max_values, requires_grad=self.quantization_parameter_learning))
+
+        w = layer.layer.weight
+        delta = qutils.calculate_delta_uniform(min_values, max_values, self.num_bits)
+        w_clipped_normed = torch.clip((w - min_values) / delta, 0, 2 ** self.num_bits - 1)
+        rest = w_clipped_normed - torch.floor(w_clipped_normed)  # rest of rounding [0, 1)
+        alpha = -torch.log((self.zeta - self.gamma) / (rest - self.gamma) - 1)  # => sigmoid(alpha) = rest
+        layer.register_parameter(f"{name}_{AUXVAR}", nn.Parameter(alpha, requires_grad=True))
+
+        # Save the quantizer parameters
+        self.add_quantizer_variable(FQ_MIN, layer.get_parameter(name+"_"+FQ_MIN), VariableGroup.QPARAMS)
+        self.add_quantizer_variable(FQ_MAX, layer.get_parameter(name+"_"+FQ_MAX), VariableGroup.QPARAMS)
+        self.add_quantizer_variable(AUXVAR, layer.get_parameter(f"{name}_{AUXVAR}"), VariableGroup.WEIGHTS)
+
+    def get_soft_targets(self) -> torch.Tensor:
+        """
+        Computes the rectified sigmoid function for the quantization target parameters.
+
+        Returns:
+            A tensor with the soft rounding targets values.
+
+        """
+        scaled_sigmoid = torch.sigmoid(self.get_quantizer_variable(AUXVAR)) * (self.zeta - self.gamma) + self.gamma
+        return torch.clip(scaled_sigmoid, min=0, max=1)
+
+    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
+
+        """
+        min_values = torch_tensor_to_numpy(self.get_quantizer_variable(FQ_MIN))
+        max_values = torch_tensor_to_numpy(self.get_quantizer_variable(FQ_MAX))
+        return {RANGE_MIN:  min_values,
+                RANGE_MAX:  max_values}
+
+    def __call__(self,
+                 inputs: nn.Parameter,
+                 training: bool) -> torch.Tensor:
+        """
+        Quantize a tensor.
+
+        Args:
+            inputs: Input tensor to quantize.
+            training: whether in training mode or not
+
+        Returns:
+            quantized tensor
+        """
+        auxvar = self.get_quantizer_variable(AUXVAR)
+        min_range = self.get_quantizer_variable(FQ_MIN)
+        max_range = self.get_quantizer_variable(FQ_MAX)
+
+        #####################################################
+        # Soft Rounding
+        #####################################################
+        aux_var = self.get_soft_targets()
+        if not training:
+            aux_var = (aux_var >= 0.5).to(auxvar.dtype)
+
+        #####################################################
+        # Quantized Input
+        #####################################################
+        q_tensor = soft_rounding_unifrom_quantizer(input_tensor=inputs,
+                                                   auxvar_tensor=aux_var,
+                                                   min_range=min_range,
+                                                   max_range=max_range,
+                                                   num_bits=self.num_bits)
+
+        return q_tensor

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

@@ -19,14 +19,14 @@ import numpy as np
 from model_compression_toolkit.core.common.defaultdict import DefaultDict
 
 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
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor, torch_tensor_to_numpy
 from model_compression_toolkit.gptq.pytorch.quantizer import quant_utils as qutils
 from model_compression_toolkit.gptq.common.gptq_constants import AUXVAR, PTQ_THRESHOLD, MAX_LSB_CHANGE
-from model_compression_toolkit.core.common.constants import THRESHOLD
+from model_compression_toolkit.constants import THRESHOLD
 from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.base_inferable_quantizer import \
     mark_quantizer

model_compression_toolkit/gptq/runner.py

@@ -15,7 +15,7 @@
 
 from typing import Callable
 
-from model_compression_toolkit import CoreConfig
+from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.statistics_correction.statistics_correction import \
     apply_statistics_correction
@@ -28,6 +28,7 @@ from model_compression_toolkit.gptq.common.gptq_training import gptq_training
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import TensorboardWriter
 from model_compression_toolkit.core.common.statistics_correction.apply_bias_correction_to_graph import \
     apply_bias_correction_to_graph
+from model_compression_toolkit.logger import Logger
 
 
 def _apply_gptq(gptq_config: GradientPTQConfigV2,
@@ -55,7 +56,7 @@ def _apply_gptq(gptq_config: GradientPTQConfigV2,
 
     """
     if gptq_config is not None and gptq_config.n_epochs > 0:
-        common.Logger.info("Using experimental Gradient Based PTQ: If you encounter an issue "
+        Logger.info("Using experimental Gradient Based PTQ: If you encounter an issue "
                            "please file a bug. To disable it, do not pass a gptq configuration.")
 
         tg_bias = gptq_training(tg,

model_compression_toolkit/{core/keras/quantization_facade.py → legacy/keras_quantization_facade.py}

@@ -15,9 +15,8 @@
 
 from typing import Callable, List, Tuple
 
-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
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.gptq import GradientPTQConfig, GradientPTQConfigV2
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
@@ -35,15 +34,15 @@ from model_compression_toolkit.ptq.runner import ptq_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.core.common.constants import FOUND_TF
+from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import TargetPlatformCapabilities
+from model_compression_toolkit.constants import FOUND_TF
 
 if FOUND_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.core.keras.keras_model_validation import KerasModelValidation
     from tensorflow.keras.models import Model
-    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 import get_target_platform_capabilities
 
@@ -81,7 +80,7 @@ if FOUND_TF:
             network_editor (List[EditRule]): List of EditRules. Each EditRule consists of a node filter and an action to change quantization settings of the filtered nodes.
             gptq_config (GradientPTQConfig): Configuration for using gptq (e.g. optimizer).
             analyze_similarity (bool): Whether to plot similarity figures within TensorBoard (when logger is enabled) or not.
-            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to. `Default Keras TPC <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/core/tpc_models/keras_tp_models/keras_default.py>`_
+            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
 
         Returns:
             A quantized model and information the user may need to handle the quantized model.
@@ -184,7 +183,7 @@ if FOUND_TF:
              network_editor (List[EditRule]): List of EditRules. Each EditRule consists of a node filter and an action to change quantization settings of the filtered nodes.
              gptq_config (GradientPTQConfig): Configuration for using GPTQ (e.g. optimizer).
              analyze_similarity (bool): Whether to plot similarity figures within TensorBoard (when logger is enabled) or not.
-             target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to. `Default Keras TPC <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/core/tpc_models/keras_tp_models/keras_default.py>`_
+             target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
 
 
          Returns:
@@ -209,13 +208,13 @@ if FOUND_TF:
              Create 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.MixedPrecisionQuantizationConfig()
+             >>> config = mct.core.MixedPrecisionQuantizationConfig()
 
              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.
 
              Pass the model, the representative dataset generator, the configuration and the target KPI to get a
              quantized model:
@@ -229,11 +228,11 @@ if FOUND_TF:
                              fw_info=fw_info).validate()
 
         if not isinstance(quant_config, MixedPrecisionQuantizationConfig):
-            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 "
                                 "MixedPrecisionQuantizationConfig. Please use keras_post_training_quantization API,"
                                 "or pass a valid mixed precision configuration.")
 
-        common.Logger.info("Using experimental mixed-precision quantization. "
+        Logger.info("Using experimental mixed-precision quantization. "
                            "If you encounter an issue please file a bug.")
 
         quantization_config, mp_config = quant_config.separate_configs()

model_compression_toolkit/{core/pytorch/quantization_facade.py → legacy/pytorch_quantization_facade.py}

@@ -14,12 +14,11 @@
 # ==============================================================================
 from typing import Callable, List, Tuple
 
-from model_compression_toolkit.core import common
-from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.core.common.constants import PYTORCH
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import PYTORCH
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig, 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.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.network_editors.actions import EditRule
@@ -34,12 +33,12 @@ from model_compression_toolkit.gptq.runner import gptq_runner
 from model_compression_toolkit.ptq.runner import ptq_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.constants import FOUND_TORCH
+from model_compression_toolkit.constants import FOUND_TORCH
 
 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.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from torch.nn import Module
 
     from model_compression_toolkit import get_target_platform_capabilities
@@ -76,7 +75,7 @@ if FOUND_TORCH:
             network_editor (List[EditRule]): List of EditRules. Each EditRule consists of a node filter and an action to change quantization settings of the filtered nodes.
             gptq_config (GradientPTQConfig): Configuration for using gptq (e.g. optimizer).
             analyze_similarity (bool): Whether to plot similarity figures within TensorBoard (when logger is enabled) or not.
-            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.
 
 
         Returns:
@@ -175,7 +174,7 @@ if FOUND_TORCH:
              network_editor (List[EditRule]): List of EditRules. Each EditRule consists of a node filter and an action to change quantization settings of the filtered nodes.
              gptq_config (GradientPTQConfig): Configuration for using GPTQ (e.g. optimizer).
              analyze_similarity (bool): Whether to plot similarity figures within TensorBoard (when logger is enabled) or not.
-             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.
 
          Returns:
              A quantized model and information the user may need to handle the quantized model.
@@ -199,13 +198,13 @@ if FOUND_TORCH:
              Create 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.MixedPrecisionQuantizationConfig()
+             >>> config = mct.core.MixedPrecisionQuantizationConfig()
 
              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 PyTorch will be affected by this value,
              while the bias will not):
 
-             >>> kpi = mct.KPI(sum(p.numel() for p in module.parameters()) * 0.75)  # About 0.75 of the model size when quantized with 8 bits.
+             >>> kpi = mct.core.KPI(sum(p.numel() for p in module.parameters()) * 0.75)  # About 0.75 of the model size when quantized with 8 bits.
 
              Pass the model, the representative dataset generator, the configuration and the target KPI to get a
              quantized model:
@@ -217,11 +216,11 @@ if FOUND_TORCH:
          """
 
         if not isinstance(quant_config, MixedPrecisionQuantizationConfig):
-            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 "
                                 "MixedPrecisionQuantizationConfig. Please use pytorch_post_training_quantization API, "
                                 "or pass a valid mixed precision configuration.")
 
-        common.Logger.info("Using experimental mixed-precision quantization. "
+        Logger.info("Using experimental mixed-precision quantization. "
                            "If you encounter an issue please file a bug.")
 
         quantization_config, mp_config = quant_config.separate_configs()

model_compression_toolkit/ptq/__init__.py

@@ -12,3 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+
+from model_compression_toolkit.ptq.pytorch.quantization_facade import pytorch_post_training_quantization_experimental
+from model_compression_toolkit.ptq.keras.quantization_facade import keras_post_training_quantization_experimental

model_compression_toolkit/ptq/keras/quantization_facade.py

@@ -15,15 +15,14 @@
 
 from typing import Callable
 
-from model_compression_toolkit import CoreConfig
-from model_compression_toolkit.core import common
+from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core.analyzer import analyzer_model_quantization
-from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.core.common.constants import TENSORFLOW, FOUND_TF
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import TENSORFLOW, FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
     MixedPrecisionQuantizationConfigV2
-from model_compression_toolkit.core.common.target_platform.targetplatform2framework import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import TargetPlatformCapabilities
 from model_compression_toolkit.core.exporter import export_model
 from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
 from model_compression_toolkit.ptq.runner import ptq_runner
@@ -33,7 +32,7 @@ if FOUND_TF:
     from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
     from model_compression_toolkit.core.keras.keras_model_validation import KerasModelValidation
     from tensorflow.keras.models import Model
-    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
@@ -93,25 +92,25 @@ if FOUND_TF:
 
             Create a MCT core config, containing the quantization configuration:
 
-            >>> config = mct.CoreConfig()
+            >>> config = mct.core.CoreConfig()
 
             If mixed precision is desired, create a 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.
             In this example we use 1 image to search mixed-precision configuration:
 
-            >>> 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.
 
             Pass the model, the representative dataset generator, the configuration and the target KPI to get a
             quantized model:
 
-            >>> quantized_model, quantization_info = mct.keras_post_training_quantization_experimental(model, repr_datagen, kpi, core_config=config)
+            >>> quantized_model, quantization_info = mct.ptq.keras_post_training_quantization_experimental(model, repr_datagen, kpi, core_config=config)
 
             For more configuration options, please take a look at our `API documentation <https://sony.github.io/model_optimization/api/api_docs/modules/mixed_precision_quantization_config.html>`_.
 
@@ -124,11 +123,11 @@ if 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)

model_compression_toolkit/ptq/pytorch/quantization_facade.py

@@ -15,11 +15,11 @@
 from typing import Callable
 
 from model_compression_toolkit.core import common
-from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.core.common.constants import PYTORCH, FOUND_TORCH
-from model_compression_toolkit.core.common.target_platform import TargetPlatformCapabilities
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import PYTORCH, FOUND_TORCH
+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 import CoreConfig
+from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
     MixedPrecisionQuantizationConfigV2
 from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
@@ -31,7 +31,7 @@ from model_compression_toolkit.core.analyzer import analyzer_model_quantization
 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.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from torch.nn import Module
     from model_compression_toolkit.exporter.model_wrapper.pytorch.builder.fully_quantized_model_builder import get_exportable_pytorch_model
     from model_compression_toolkit import get_target_platform_capabilities
@@ -88,18 +88,18 @@ if FOUND_TORCH:
             Set number of clibration iterations to 1:
 
             >>> import model_compression_toolkit as mct
-            >>> quantized_module, quantization_info = mct.pytorch_post_training_quantization_experimental(module, repr_datagen)
+            >>> quantized_module, quantization_info = mct.ptq.pytorch_post_training_quantization_experimental(module, repr_datagen)
 
         """
 
         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 "
                                     "pytorch_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)

model_compression_toolkit/qat/__init__.py

@@ -12,3 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+from model_compression_toolkit.qat.common.qat_config import QATConfig, TrainingMethod
+
+from model_compression_toolkit.qat.keras.quantization_facade import keras_quantization_aware_training_init, keras_quantization_aware_training_finalize
+from model_compression_toolkit.qat.pytorch.quantization_facade import pytorch_quantization_aware_training_init, pytorch_quantization_aware_training_finalize

model_compression_toolkit/qat/common/__init__.py

@@ -12,5 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
-from model_compression_toolkit.qat.common.constants import THRESHOLD_TENSOR, WEIGHTS_QUANTIZATION_PARAMS
+from model_compression_toolkit.quantizers_infrastructure.constants import THRESHOLD_TENSOR, WEIGHTS_QUANTIZATION_PARAMS

model_compression_toolkit/qat/common/qat_config.py

@@ -17,6 +17,8 @@ from typing import Dict
 from enum import Enum
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+from model_compression_toolkit.logger import Logger
+
 
 def _is_qat_applicable(node: common.BaseNode,
                        fw_info: FrameworkInfo) -> bool:
@@ -31,7 +33,7 @@ def _is_qat_applicable(node: common.BaseNode,
     """
 
     if node.is_weights_quantization_enabled() and not fw_info.is_kernel_op(node.type):
-        common.Logger.error("QAT Error: Quantizing a node without a kernel isn't supported")
+        Logger.error("QAT Error: Quantizing a node without a kernel isn't supported")
     return node.is_weights_quantization_enabled() or node.is_activation_quantization_enabled()
 
 
@@ -40,8 +42,10 @@ class TrainingMethod(Enum):
     An enum for selecting a QAT training method
 
     STE - Standard straight-through estimator. Includes PowerOfTwo, symmetric & uniform quantizers
+    DQA -  DNN Quantization with Attention. Includes a smooth quantization introduces by DQA method
     """
     STE = "STE",
+    DQA = "DQA"
 
 
 class QATConfig: