mct-nightly 1.8.0.20052023.post401__py3-none-any.whl → 1.8.0.20230610.post356__py3-none-any.whl

model_compression_toolkit/gptq/pytorch/quantizer/quantization_builder.py

@@ -21,15 +21,13 @@ from model_compression_toolkit.exporter.model_wrapper.pytorch.builder.node_to_qu
     get_activation_inferable_quantizer_kwargs
 from model_compression_toolkit.gptq.pytorch.quantizer.base_pytorch_gptq_quantizer import \
     BasePytorchGPTQTrainableQuantizer
-from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.get_quantizers import \
-    get_inferable_quantizer_class
-from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.pytorch.quantizers import \
-    BasePyTorchInferableQuantizer
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.get_quantizer_config import \
+from mct_quantizers import QuantizationTarget
+from mct_quantizers.common.get_quantizers import get_inferable_quantizer_class
+from mct_quantizers.pytorch.quantizers import BasePyTorchInferableQuantizer
+from model_compression_toolkit.trainable_infrastructure.common.get_quantizer_config import \
     get_trainable_quantizer_weights_config
 from model_compression_toolkit.qat.pytorch.quantizer.base_pytorch_qat_quantizer import BasePytorchQATTrainableQuantizer
-from model_compression_toolkit.quantizers_infrastructure import QuantizationTarget
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.get_quantizers import \
+from model_compression_toolkit.trainable_infrastructure.common.get_quantizers import \
     get_trainable_quantizer_class
 
 

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

@@ -21,7 +21,7 @@ from torch import nn
 from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
-from model_compression_toolkit.quantizers_infrastructure import PytorchQuantizationWrapper
+from mct_quantizers import PytorchQuantizationWrapper
 
 
 class LinearTempDecay:

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

@@ -18,8 +18,8 @@ from typing import Dict
 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.target_platform_capabilities.target_platform import QuantizationMethod
+from mct_quantizers import QuantizationTarget, PytorchQuantizationWrapper
 from model_compression_toolkit.gptq.common.gptq_config import RoundingType
 from model_compression_toolkit.gptq.pytorch.quantizer.base_pytorch_gptq_quantizer import \
     BasePytorchGPTQTrainableQuantizer
@@ -28,11 +28,11 @@ from model_compression_toolkit.gptq.pytorch.quantizer import quant_utils as quti
 from model_compression_toolkit.gptq.common.gptq_constants import PTQ_THRESHOLD, SCALE_PTQ, \
     SOFT_ROUNDING_GAMMA, SOFT_ROUNDING_ZETA, AUXVAR
 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 \
+from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig
+from mct_quantizers import mark_quantizer
+from model_compression_toolkit.trainable_infrastructure.common.quant_utils import \
     get_threshold_reshape_shape
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
+from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
 def soft_rounding_symmetric_quantizer(input_tensor: torch.Tensor,
@@ -68,7 +68,7 @@ def soft_rounding_symmetric_quantizer(input_tensor: torch.Tensor,
                                    max_val=int_threshold - 1)
 
 
-@mark_quantizer(quantization_target=qi.QuantizationTarget.Weights,
+@mark_quantizer(quantization_target=QuantizationTarget.Weights,
                 quantization_method=[QuantizationMethod.POWER_OF_TWO, QuantizationMethod.SYMMETRIC],
                 quantizer_type=RoundingType.SoftQuantizer)
 class SymmetricSoftRoundingGPTQ(BasePytorchGPTQTrainableQuantizer):
@@ -110,7 +110,7 @@ class SymmetricSoftRoundingGPTQ(BasePytorchGPTQTrainableQuantizer):
     def initialize_quantization(self,
                                 tensor_shape: torch.Size,
                                 name: str,
-                                layer: qi.PytorchQuantizationWrapper):
+                                layer: PytorchQuantizationWrapper):
         """
         Add quantizer parameters to the quantizer parameters dictionary
 

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

@@ -17,9 +17,9 @@ 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.trainable_infrastructure.common.constants import FQ_MIN, FQ_MAX
 from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
+from mct_quantizers import QuantizationTarget, PytorchQuantizationWrapper
 from model_compression_toolkit.gptq.common.gptq_config import RoundingType
 from model_compression_toolkit.gptq.pytorch.quantizer.base_pytorch_gptq_quantizer import \
     BasePytorchGPTQTrainableQuantizer
@@ -27,10 +27,9 @@ 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 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 \
+from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig
+from mct_quantizers import mark_quantizer
+from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import \
     VariableGroup
 from model_compression_toolkit.constants import RANGE_MAX, RANGE_MIN
 
@@ -63,7 +62,7 @@ def soft_rounding_unifrom_quantizer(input_tensor: torch.Tensor,
                                    max_val=2 ** num_bits - 1) + min_range
 
 
-@mark_quantizer(quantization_target=qi.QuantizationTarget.Weights,
+@mark_quantizer(quantization_target=QuantizationTarget.Weights,
                 quantization_method=[QuantizationMethod.UNIFORM],
                 quantizer_type=RoundingType.SoftQuantizer)
 class UniformSoftRoundingGPTQ(BasePytorchGPTQTrainableQuantizer):
@@ -100,7 +99,7 @@ class UniformSoftRoundingGPTQ(BasePytorchGPTQTrainableQuantizer):
     def initialize_quantization(self,
                                 tensor_shape: torch.Size,
                                 name: str,
-                                layer: qi.PytorchQuantizationWrapper):
+                                layer: PytorchQuantizationWrapper):
         """
         Add quantizer parameters to the quantizer parameters dictionary
 

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

@@ -18,8 +18,8 @@ from typing import Dict
 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.target_platform_capabilities.target_platform import QuantizationMethod
+from mct_quantizers import QuantizationTarget, PytorchQuantizationWrapper
 from model_compression_toolkit.gptq.common.gptq_config import RoundingType
 from model_compression_toolkit.gptq.pytorch.quantizer.base_pytorch_gptq_quantizer import \
     BasePytorchGPTQTrainableQuantizer
@@ -27,11 +27,10 @@ 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 AUXVAR, PTQ_THRESHOLD, MAX_LSB_CHANGE
 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
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.quant_utils import \
+from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig
+from mct_quantizers import mark_quantizer
+from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
+from model_compression_toolkit.trainable_infrastructure.common.quant_utils import \
     get_threshold_reshape_shape
 
 
@@ -75,7 +74,7 @@ def pertubation_symmetric_quantizer(input_tensor: torch.Tensor,
     return delta * qutils.ste_clip(tensor_q, max_val=max_int, min_val=min_int)
 
 
-@mark_quantizer(quantization_target=qi.QuantizationTarget.Weights,
+@mark_quantizer(quantization_target=QuantizationTarget.Weights,
                 quantization_method=[QuantizationMethod.POWER_OF_TWO, QuantizationMethod.SYMMETRIC],
                 quantizer_type=RoundingType.STE)
 class STEWeightGPTQQuantizer(BasePytorchGPTQTrainableQuantizer):
@@ -109,7 +108,7 @@ class STEWeightGPTQQuantizer(BasePytorchGPTQTrainableQuantizer):
     def initialize_quantization(self,
                                 tensor_shape: torch.Size,
                                 name: str,
-                                layer: qi.PytorchQuantizationWrapper):
+                                layer: PytorchQuantizationWrapper):
         """
         Add quantizer parameters to the quantizer parameters dictionary
 

model_compression_toolkit/qat/common/__init__.py

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

model_compression_toolkit/qat/common/qat_config.py

@@ -20,8 +20,8 @@ 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:
+def is_qat_applicable(node: common.BaseNode,
+                      fw_info: FrameworkInfo) -> bool:
     """
     A function for deciding if a layer should be fine-tuned during QAT
     Args:

model_compression_toolkit/qat/keras/quantization_facade.py

@@ -22,7 +22,7 @@ from model_compression_toolkit.constants import 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.quantizers_infrastructure import ActivationQuantizationHolder
+from mct_quantizers import KerasActivationQuantizationHolder, KerasQuantizationWrapper
 from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import TargetPlatformCapabilities
 from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
 from model_compression_toolkit.ptq.runner import ptq_runner
@@ -40,20 +40,18 @@ if FOUND_TF:
     from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder
 
     from model_compression_toolkit import get_target_platform_capabilities
-    from model_compression_toolkit import quantizers_infrastructure as qi
 
     from model_compression_toolkit import get_target_platform_capabilities
     from model_compression_toolkit.core import common
     from model_compression_toolkit.core.common import BaseNode
     from model_compression_toolkit.constants import TENSORFLOW
     from model_compression_toolkit.core.common.framework_info import FrameworkInfo
-    from model_compression_toolkit.qat.common.qat_config import _is_qat_applicable
+    from model_compression_toolkit.qat.common.qat_config import is_qat_applicable
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
     from model_compression_toolkit.qat.keras.quantizer.quantization_builder import quantization_builder, \
     get_activation_quantizer_holder
     from model_compression_toolkit.qat.common.qat_config import QATConfig
-    from model_compression_toolkit import quantizers_infrastructure as qi
 
     DEFAULT_KERAS_TPC = get_target_platform_capabilities(TENSORFLOW, DEFAULT_TP_MODEL)
 
@@ -71,9 +69,12 @@ if FOUND_TF:
         Returns: Wrapped layer
 
         """
-        if _is_qat_applicable(n, DEFAULT_KERAS_INFO):
-            weights_quantizers, activation_quantizers = quantization_builder(n, qat_config, DEFAULT_KERAS_INFO)
-            return qi.KerasQuantizationWrapper(layer, weights_quantizers)
+        if is_qat_applicable(n, DEFAULT_KERAS_INFO):
+            weights_quantizers, _ = quantization_builder(n,
+                                                         qat_config,
+                                                         DEFAULT_KERAS_INFO)
+            if len(weights_quantizers) > 0:
+                return KerasQuantizationWrapper(layer, weights_quantizers)
         return layer
 
 
@@ -255,8 +256,17 @@ if FOUND_TF:
 
          """
         def _export(layer):
-            if isinstance(layer, (qi.KerasQuantizationWrapper, ActivationQuantizationHolder)):
+            if isinstance(layer, KerasQuantizationWrapper):
                 layer.convert_to_inferable_quantizers()
+            # In the KerasActivationQuantizationHolder case - converting the quantizers only
+            # is not enough. We need to create a new layer with inferable quantizers. The reason for that
+            # is that if we only convert the quantizers, the layer will have some weights (such as min, max,
+            # threshold) that do not match the configuration, thus loading such a model will fail.
+            # To overcome this, the convert_to_inferable_quantizers of KerasActivationQuantizationHolder
+            # creates a new layer from its new configuration after converting the trainable quantizer
+            # to an inferable quantizer.
+            elif isinstance(layer, KerasActivationQuantizationHolder):
+                layer = layer.convert_to_inferable_quantizers()
             return layer
 
         # clone each layer in the model and apply _export to layers with TrainableQuantizeWrappers

model_compression_toolkit/qat/keras/quantizer/base_keras_qat_quantizer.py

@@ -17,7 +17,7 @@ from typing import Union
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import FOUND_TF
 
-from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig, \
+from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig, \
     TrainableQuantizerActivationConfig, BaseKerasTrainableQuantizer
 
 if FOUND_TF:

model_compression_toolkit/qat/keras/quantizer/quantization_builder.py

@@ -12,34 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Tuple, Dict, List, Union, Callable
+from typing import Tuple, Dict, List, Callable
 
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.qat.common.qat_config import QATConfig, _is_qat_applicable
+from model_compression_toolkit.qat.common.qat_config import QATConfig
 from model_compression_toolkit.qat.keras.quantizer.base_keras_qat_quantizer import BaseKerasQATTrainableQuantizer
-from model_compression_toolkit.quantizers_infrastructure import QuantizationTarget, ActivationQuantizationHolder
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.get_quantizer_config import \
+from mct_quantizers import QuantizationTarget, KerasActivationQuantizationHolder
+from model_compression_toolkit.trainable_infrastructure.common.get_quantizer_config import \
     get_trainable_quantizer_weights_config, get_trainable_quantizer_activation_config, \
     get_trainable_quantizer_quantization_candidates
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.get_quantizers import \
+from model_compression_toolkit.trainable_infrastructure.common.get_quantizers import \
     get_trainable_quantizer_class
 
 
 def get_activation_quantizer_holder(n: common.BaseNode,
-                                    qat_config: QATConfig) -> Union[None, Callable]:
+                                    qat_config: QATConfig) -> Callable:
     """
-    Retrieve a ActivationQuantizationHolder layer to use for activation quantization for a node.
+    Retrieve a KerasActivationQuantizationHolder layer to use for activation quantization for a node.
     If the layer is not supposed to be wrapped with activation quantizers - return None.
 
     Args:
-        n: Node to get ActivationQuantizationHolder to attach in its output.
+        n: Node to get KerasActivationQuantizationHolder to attach in its output.
         qat_config: Configuration of QAT (such as training methods for example).
 
     Returns:
-        A ActivationQuantizationHolder layer for the node activation quantization.
+        A KerasActivationQuantizationHolder layer for the node activation quantization.
     """
     _, activation_quantizers = quantization_builder(n,
                                                     qat_config,
@@ -49,8 +49,8 @@ def get_activation_quantizer_holder(n: common.BaseNode,
     # thus we make sure this is the only possible case (unless it's a node with no activation
     # quantization, which in this case has an empty list).
     if len(activation_quantizers) == 1:
-        return ActivationQuantizationHolder(activation_quantizers[0])
-    Logger.error(f'ActivationQuantizationHolder supports a single quantizer but {len(activation_quantizers)} quantizers were found for node {n}')
+        return KerasActivationQuantizationHolder(activation_quantizers[0])
+    Logger.error(f'KerasActivationQuantizationHolder supports a single quantizer but {len(activation_quantizers)} quantizers were found for node {n}')
 
 
 def quantization_builder(n: common.BaseNode,

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

@@ -19,25 +19,24 @@ import numpy as np
 import tensorflow as tf
 from tensorflow.python.framework.tensor_shape import TensorShape
 from model_compression_toolkit.constants import SIGNED
-from model_compression_toolkit.quantizers_infrastructure.constants import FQ_MIN, FQ_MAX
+from model_compression_toolkit.trainable_infrastructure.common.constants import FQ_MIN, FQ_MAX
 
 from model_compression_toolkit.qat import TrainingMethod
 
 from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
+from mct_quantizers import QuantizationTarget, mark_quantizer, KerasQuantizationWrapper
 from model_compression_toolkit.qat.common import THRESHOLD_TENSOR
-from model_compression_toolkit import quantizers_infrastructure as qi, constants as C
+from model_compression_toolkit import constants as C
 
 from model_compression_toolkit.qat.keras.quantizer.base_keras_qat_quantizer import BaseKerasQATTrainableQuantizer
-from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig, \
+from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig, \
     TrainableQuantizerActivationConfig
-from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.base_inferable_quantizer import mark_quantizer
-from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers import \
-    WeightsPOTInferableQuantizer, WeightsSymmetricInferableQuantizer, ActivationPOTInferableQuantizer, \
-    ActivationSymmetricInferableQuantizer
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
+from mct_quantizers.keras.quantizers import WeightsPOTInferableQuantizer, WeightsSymmetricInferableQuantizer, \
+    ActivationPOTInferableQuantizer, ActivationSymmetricInferableQuantizer
+from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
-@mark_quantizer(quantization_target=qi.QuantizationTarget.Weights,
+@mark_quantizer(quantization_target=QuantizationTarget.Weights,
                 quantization_method=[QuantizationMethod.POWER_OF_TWO, QuantizationMethod.SYMMETRIC],
                 quantizer_type=TrainingMethod.STE)
 class STEWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
@@ -84,7 +83,7 @@ class STEWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
     def initialize_quantization(self,
                                 tensor_shape: TensorShape,
                                 name: str,
-                                layer: qi.KerasQuantizationWrapper):
+                                layer: KerasQuantizationWrapper):
         """
         Add quantizer parameters to the quantizer parameters dictionary
 
@@ -171,7 +170,7 @@ class STEWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
                                                       input_rank=len(self.threshold_shape))
 
 
-@mark_quantizer(quantization_target=qi.QuantizationTarget.Activation,
+@mark_quantizer(quantization_target=QuantizationTarget.Activation,
                 quantization_method=[QuantizationMethod.POWER_OF_TWO, QuantizationMethod.SYMMETRIC],
                 quantizer_type=TrainingMethod.STE)
 class STEActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
@@ -206,7 +205,7 @@ class STEActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
     def initialize_quantization(self,
                                 tensor_shape: TensorShape,
                                 name: str,
-                                layer: qi.KerasQuantizationWrapper):
+                                layer: KerasQuantizationWrapper):
         """
         Add quantizer parameters to the quantizer parameters dictionary
 

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

@@ -16,25 +16,24 @@ import numpy as np
 import tensorflow as tf
 from tensorflow.python.framework.tensor_shape import TensorShape
 from model_compression_toolkit.constants import RANGE_MIN, RANGE_MAX
-from model_compression_toolkit.quantizers_infrastructure.constants import FQ_MIN, FQ_MAX
+from model_compression_toolkit.trainable_infrastructure.common.constants import FQ_MIN, FQ_MAX
 from model_compression_toolkit.qat import TrainingMethod
-from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
+
+from mct_quantizers import mark_quantizer, QuantizationMethod, QuantizationTarget, KerasQuantizationWrapper
+from mct_quantizers.keras.quantizers import \
+    BaseKerasInferableQuantizer, WeightsUniformInferableQuantizer, ActivationUniformInferableQuantizer
 
 from model_compression_toolkit.qat.keras.quantizer.quant_utils import adjust_range_to_include_zero
 from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import fix_range_to_include_zero
-from model_compression_toolkit import quantizers_infrastructure as qi, constants as C
+from model_compression_toolkit import constants as C
 
 from model_compression_toolkit.qat.keras.quantizer.base_keras_qat_quantizer import BaseKerasQATTrainableQuantizer
-from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig, \
+from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig, \
     TrainableQuantizerActivationConfig
-from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.base_inferable_quantizer import \
-    mark_quantizer
-from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers import \
-    BaseKerasInferableQuantizer, WeightsUniformInferableQuantizer, ActivationUniformInferableQuantizer
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
+from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
-@mark_quantizer(quantization_target=qi.QuantizationTarget.Weights,
+@mark_quantizer(quantization_target=QuantizationTarget.Weights,
                 quantization_method=[QuantizationMethod.UNIFORM],
                 quantizer_type=TrainingMethod.STE)
 class STEUniformWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
@@ -73,7 +72,7 @@ class STEUniformWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
     def initialize_quantization(self,
                                 tensor_shape: TensorShape,
                                 name: str,
-                                layer: qi.KerasQuantizationWrapper):
+                                layer: KerasQuantizationWrapper):
         """
         Add quantizer parameters to the quantizer parameters dictionary
 
@@ -148,7 +147,7 @@ class STEUniformWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
                                                 input_rank=len(self.min_max_shape))
 
 
-@mark_quantizer(quantization_target=qi.QuantizationTarget.Activation,
+@mark_quantizer(quantization_target=QuantizationTarget.Activation,
                 quantization_method=[QuantizationMethod.UNIFORM],
                 quantizer_type=TrainingMethod.STE)
 class STEUniformActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
@@ -173,7 +172,7 @@ class STEUniformActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
     def initialize_quantization(self,
                                 tensor_shape: TensorShape,
                                 name: str,
-                                layer: qi.KerasQuantizationWrapper):
+                                layer: KerasQuantizationWrapper):
         """
         Add quantizer parameters to the quantizer parameters dictionary
 

model_compression_toolkit/qat/pytorch/quantization_facade.py

@@ -25,28 +25,32 @@ from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 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.target_platform_capabilities.target_platform.targetplatform2framework import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import \
+    TargetPlatformCapabilities
 from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
 from model_compression_toolkit.ptq.runner import ptq_runner
 
-
 if FOUND_TORCH:
     import torch.nn as nn
     from torch.nn import Module
+    from mct_quantizers import PytorchActivationQuantizationHolder
     from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from model_compression_toolkit.core.pytorch.pytorch_implementation import PytorchImplementation
-    from model_compression_toolkit.qat.common.qat_config import _is_qat_applicable
+    from model_compression_toolkit.qat.common.qat_config import is_qat_applicable
     from model_compression_toolkit.core.pytorch.back2framework.pytorch_model_builder import PyTorchModelBuilder
-    from model_compression_toolkit.quantizers_infrastructure import PytorchQuantizationWrapper
-    from model_compression_toolkit import quantizers_infrastructure as qi
+    from mct_quantizers import PytorchQuantizationWrapper
     from model_compression_toolkit import get_target_platform_capabilities
     from model_compression_toolkit.qat.common.qat_config import QATConfig
+    from model_compression_toolkit.qat.pytorch.quantizer.quantization_builder import get_activation_quantizer_holder
     from model_compression_toolkit.qat.pytorch.quantizer.quantization_builder import quantization_builder
+
     DEFAULT_PYTORCH_TPC = get_target_platform_capabilities(PYTORCH, DEFAULT_TP_MODEL)
 
 
-    def qat_wrapper(n: common.BaseNode, module: nn.Module, qat_config: QATConfig):
+    def qat_wrapper(n: common.BaseNode,
+                    module: nn.Module,
+                    qat_config: QATConfig):
         """
         A function which takes a computational graph node and a pytorch module and perform the quantization wrapping
         Args:
@@ -56,11 +60,11 @@ if FOUND_TORCH:
         Returns: Wrapped layer
 
         """
-        if _is_qat_applicable(n, DEFAULT_PYTORCH_INFO):
-            weights_quantizers, activation_quantizers = quantization_builder(n, qat_config, DEFAULT_PYTORCH_INFO)
-            return qi.PytorchQuantizationWrapper(module, weights_quantizers, activation_quantizers)
-        else:
-            return module
+        if is_qat_applicable(n, DEFAULT_PYTORCH_INFO):
+            weights_quantizers, _ = quantization_builder(n, qat_config, DEFAULT_PYTORCH_INFO)
+            if len(weights_quantizers) > 0:
+                return PytorchQuantizationWrapper(module, weights_quantizers)
+        return module
 
 
     def pytorch_quantization_aware_training_init(in_model: Module,
@@ -135,11 +139,11 @@ if FOUND_TORCH:
         if core_config.mixed_precision_enable:
             if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfigV2):
                 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.")
+                             "MixedPrecisionQuantizationConfigV2. Please use pytorch_post_training_quantization API,"
+                             "or pass a valid mixed precision configuration.")
 
             Logger.info("Using experimental mixed-precision quantization. "
-                               "If you encounter an issue please file a bug.")
+                        "If you encounter an issue please file a bug.")
 
         tb_w = _init_tensorboard_writer(fw_info)
 
@@ -158,12 +162,18 @@ if FOUND_TORCH:
 
         _qat_wrapper = partial(qat_wrapper, qat_config=qat_config)
 
-        qat_model, user_info = PyTorchModelBuilder(graph=tg, fw_info=fw_info, wrapper=_qat_wrapper).build_model()
+        qat_model, user_info = PyTorchModelBuilder(graph=tg,
+                                                   fw_info=fw_info,
+                                                   wrapper=_qat_wrapper,
+                                                   get_activation_quantizer_holder_fn=partial(
+                                                       get_activation_quantizer_holder,
+                                                       qat_config=qat_config)).build_model()
 
         user_info.mixed_precision_cfg = bit_widths_config
 
         return qat_model, user_info
 
+
     def pytorch_quantization_aware_training_finalize(in_model: Module):
         """
          Convert a model fine-tuned by the user to a network with QuantizeWrappers containing
@@ -207,7 +217,7 @@ if FOUND_TORCH:
          """
         exported_model = copy.deepcopy(in_model)
         for _, layer in exported_model.named_children():
-            if isinstance(layer, PytorchQuantizationWrapper):
+            if isinstance(layer, (PytorchQuantizationWrapper, PytorchActivationQuantizationHolder)):
                 layer.convert_to_inferable_quantizers()
 
         return exported_model
@@ -221,6 +231,7 @@ else:
                         'when using pytorch_quantization_aware_training_init. '
                         'Could not find the torch package.')  # pragma: no cover
 
+
     def pytorch_quantization_aware_training_finalize(*args, **kwargs):
         Logger.critical('Installing Pytorch is mandatory '
                         'when using pytorch_quantization_aware_training_finalize. '

model_compression_toolkit/qat/pytorch/quantizer/base_pytorch_qat_quantizer.py

@@ -17,9 +17,9 @@ from typing import Union
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import FOUND_TORCH
 
-from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig, \
+from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig, \
     TrainableQuantizerActivationConfig
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.pytorch.base_pytorch_quantizer import \
+from model_compression_toolkit.trainable_infrastructure.pytorch.base_pytorch_quantizer import \
     BasePytorchTrainableQuantizer
 
 if FOUND_TORCH:

model_compression_toolkit/qat/pytorch/quantizer/quantization_builder.py

@@ -12,19 +12,46 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import List, Dict, Tuple
+from typing import List, Dict, Tuple, Callable
+
+from mct_quantizers import PytorchActivationQuantizationHolder, QuantizationTarget
 
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.qat.common.qat_config import QATConfig
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.get_quantizer_config import \
+from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.trainable_infrastructure.common.get_quantizer_config import \
     get_trainable_quantizer_quantization_candidates, get_trainable_quantizer_weights_config, \
     get_trainable_quantizer_activation_config
 from model_compression_toolkit.qat.pytorch.quantizer.base_pytorch_qat_quantizer import BasePytorchQATTrainableQuantizer
-from model_compression_toolkit.quantizers_infrastructure import QuantizationTarget
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.get_quantizers import \
+from model_compression_toolkit.trainable_infrastructure.common.get_quantizers import \
     get_trainable_quantizer_class
 
+def get_activation_quantizer_holder(n: common.BaseNode,
+                                    qat_config: QATConfig) -> Callable:
+    """
+    Retrieve a ActivationQuantizationHolder layer to use for activation quantization for a node.
+    If the layer is not supposed to be wrapped with activation quantizers - return None.
+
+    Args:
+        n: Node for which to retrieve anActivationQuantizationHolder to attach to its output.
+        qat_config: QAT configuration (for example, training methods).
+
+    Returns:
+        A ActivationQuantizationHolder layer for the node's activation quantization.
+    """
+    _, activation_quantizers = quantization_builder(n,
+                                                    qat_config,
+                                                    DEFAULT_PYTORCH_INFO)
+
+    # Holder by definition uses a single quantizer for the activation quantization
+    # thus we make sure this is the only possible case (unless it's a node with no activation
+    # quantization, which in this case has an empty list).
+    if len(activation_quantizers) == 1:
+        return PytorchActivationQuantizationHolder(activation_quantizers[0])
+    Logger.error(f'ActivationQuantizationHolder supports a single quantizer but {len(activation_quantizers)} quantizers were found for node {n}')
+
 
 def quantization_builder(n: common.BaseNode,
                          qat_config: QATConfig,