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

model_compression_toolkit/exporter/model_wrapper/keras/validate_layer.py

@@ -14,16 +14,14 @@
 # ==============================================================================
 from typing import Any
 
-
-from model_compression_toolkit.logger import Logger
+from mct_quantizers import BaseInferableQuantizer, KerasActivationQuantizationHolder
 from model_compression_toolkit.constants import FOUND_TF
-
-from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.base_inferable_quantizer import BaseInferableQuantizer
-
+from model_compression_toolkit.logger import Logger
 
 if FOUND_TF:
+    from keras.engine.base_layer import Layer
     from keras.engine.input_layer import InputLayer
-    from model_compression_toolkit.quantizers_infrastructure import KerasQuantizationWrapper
+    from mct_quantizers import KerasQuantizationWrapper
 
     def is_keras_layer_exportable(layer: Any) -> bool:
         """
@@ -39,40 +37,34 @@ if FOUND_TF:
         if isinstance(layer, InputLayer):
             return True
 
-        valid_layer = isinstance(layer, KerasQuantizationWrapper)
+        valid_layer = isinstance(layer, Layer)
         if not valid_layer:
             Logger.error(
-                f'Exportable layer must be wrapped using KerasQuantizationWrapper, but layer {layer.name} is of type '
+                f'Exportable layer must be a Keras layer, but layer {layer.name} is of type '
                 f'{type(layer)}') # pragma: no cover
 
-        valid_weights_quantizers = isinstance(layer.weights_quantizers, dict)
-        if not valid_weights_quantizers:
-            Logger.error(
-                f'KerasQuantizationWrapper must have a weights_quantizers but has a '
-                f'{type(layer.weights_quantizers)} object') # pragma: no cover
-
-        for _, weights_quantizer in layer.weights_quantizers.items():
-            if not isinstance(weights_quantizer, BaseInferableQuantizer):
+        if isinstance(layer, KerasQuantizationWrapper):
+            valid_weights_quantizers = isinstance(layer.weights_quantizers, dict)
+            if not valid_weights_quantizers:
                 Logger.error(
-                    f'weights_quantizer must be a BaseInferableQuantizer object but has a '
-                    f'{type(weights_quantizer)} object')  # pragma: no cover
+                    f'KerasQuantizationWrapper must have a weights_quantizers but has a '
+                    f'{type(layer.weights_quantizers)} object') # pragma: no cover
 
-        valid_activation_quantizers = isinstance(layer.activation_quantizers, list)
-        if not valid_activation_quantizers:
-            Logger.error(
-                f'KerasQuantizationWrapper must have a activation_quantizers list but has a '
-                f'{type(layer.activation_quantizers)} object') # pragma: no cover
+            if len(layer.weights_quantizers) == 0:
+                Logger.error(f'KerasQuantizationWrapper must have at least one weight quantizer, but found {len(layer.weights_quantizers)} quantizers. If layer is not quantized it should be a Keras layer.')
 
-        for activation_quantizers in layer.activation_quantizers:
-            if not isinstance(activation_quantizers, BaseInferableQuantizer):
-                Logger.error(
-                    f'activation_quantizers must be a BaseInferableQuantizer object but has a '
-                    f'{type(activation_quantizers)} object')  # pragma: no cover
+            for _, weights_quantizer in layer.weights_quantizers.items():
+                if not isinstance(weights_quantizer, BaseInferableQuantizer):
+                    Logger.error(
+                        f'weights_quantizer must be a BaseInferableQuantizer object but has a '
+                        f'{type(weights_quantizer)} object')  # pragma: no cover
 
-        quantizers = layer.activation_quantizers + list(layer.weights_quantizers.values())
-        is_valid_quantizers = all([isinstance(x, BaseInferableQuantizer) for x in quantizers])
-        if not is_valid_quantizers:
-            Logger.error(f'Found a quantizer that is not of type BaseInferableQuantizer') # pragma: no cover
+        if isinstance(layer, KerasActivationQuantizationHolder):
+            if not isinstance(layer.activation_holder_quantizer, BaseInferableQuantizer):
+                Logger.error(
+                    f'activation quantizer in KerasActivationQuantizationHolder'
+                    f' must be a BaseInferableQuantizer object but has a '
+                    f'{type(layer.activation_holder_quantizer)} object')  # pragma: no cover
 
         return True
 else:

model_compression_toolkit/exporter/model_wrapper/pytorch/builder/fully_quantized_model_builder.py

@@ -13,20 +13,23 @@
 # limitations under the License.
 # ==============================================================================
 
-
-from model_compression_toolkit import quantizers_infrastructure as qi
+from typing import Union, Callable
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.constants import FOUND_TORCH
 from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.core.common import BaseNode
 
 if FOUND_TORCH:
     import torch
+    from mct_quantizers import PytorchQuantizationWrapper, PytorchActivationQuantizationHolder
     from model_compression_toolkit.core.pytorch.back2framework.pytorch_model_builder import PyTorchModelBuilder
     from model_compression_toolkit.exporter.model_wrapper.pytorch.builder.node_to_quantizers import \
         get_quantization_quantizers
 
-    def fully_quantized_wrapper(node: common.BaseNode, module: torch.nn.Module) -> qi.PytorchQuantizationWrapper:
+
+    def fully_quantized_wrapper(node: common.BaseNode,
+                                module: torch.nn.Module) -> Union[torch.nn.Module,PytorchQuantizationWrapper]:
         """
         A function which takes a computational graph node and a pytorch module and
         perform the quantization wrapping
@@ -34,14 +37,32 @@ if FOUND_TORCH:
         Args:
             node: A node of mct graph.
             module: A Pytorch module
-
         Returns: Wrapped layer
 
         """
-        weight_quantizers, activation_quantizers = get_quantization_quantizers(node)
-        wrapped_layer = qi.PytorchQuantizationWrapper(module, weight_quantizers, activation_quantizers)
-        return wrapped_layer
+        weight_quantizers, _ = get_quantization_quantizers(node)
+        if len(weight_quantizers) > 0:
+            return PytorchQuantizationWrapper(module, weight_quantizers)
+        return module
 
+    def get_activation_quantizer_holder(node: BaseNode) -> Callable:
+        """
+        Retrieve a PytorchActivationQuantizationHolder layer to use for activation quantization of a node.
+        If the layer is not supposed to be wrapped with an activation quantizer - return None.
+        Args:
+            node: Node to attach a PytorchActivationQuantizationHolder to its output.
+        Returns:
+            A PytorchActivationQuantizationHolder module for the node's activation quantization.
+        """
+        _, activation_quantizers = get_quantization_quantizers(node)
+        # 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 we no activation
+        # quantization, which in this case has an empty list).
+        if len(activation_quantizers) == 1:
+            return PytorchActivationQuantizationHolder(activation_quantizers[0])
+        Logger.error(
+            f'PytorchActivationQuantizationHolder supports a single quantizer but {len(activation_quantizers)} quantizers '
+            f'were found for node {node}')
 
     def get_exportable_pytorch_model(graph: Graph):
         """
@@ -54,7 +75,9 @@ if FOUND_TORCH:
             Fully quantized PyTorch model.
         """
         return PyTorchModelBuilder(graph=graph,
-                                   wrapper=fully_quantized_wrapper).build_model()
+                                   wrapper=fully_quantized_wrapper,
+                                   get_activation_quantizer_holder_fn=get_activation_quantizer_holder).build_model()
+
 else:
     def get_exportable_pytorch_model(*args, **kwargs):  # pragma: no cover
         Logger.error('Installing torch is mandatory '

model_compression_toolkit/exporter/model_wrapper/pytorch/builder/node_to_quantizer.py

@@ -20,11 +20,11 @@ from model_compression_toolkit.constants import THRESHOLD, SIGNED, RANGE_MIN, RA
     SCALE_PER_CHANNEL, CLUSTER_CENTERS
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
-from model_compression_toolkit.quantizers_infrastructure import QuantizationTarget
-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 \
-    constants as qi_inferable_quantizers_constants, BasePyTorchInferableQuantizer
+from mct_quantizers import QuantizationTarget
+from mct_quantizers.common.get_quantizers import get_inferable_quantizer_class
+from mct_quantizers import \
+    constants as qi_inferable_quantizers_constants
+from mct_quantizers.pytorch.quantizers import BasePyTorchInferableQuantizer
 import numpy as np
 
 

model_compression_toolkit/exporter/model_wrapper/pytorch/validate_layer.py

@@ -17,10 +17,13 @@ from typing import Any
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import FOUND_TORCH
 
+
 if FOUND_TORCH:
-    from model_compression_toolkit.quantizers_infrastructure import PytorchQuantizationWrapper
-    from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.pytorch.quantizers import \
-        BasePyTorchInferableQuantizer
+    import torch.nn as nn
+    from mct_quantizers import PytorchQuantizationWrapper
+    from mct_quantizers.pytorch.quantizers import BasePyTorchInferableQuantizer
+    from mct_quantizers.pytorch.activation_quantization_holder import PytorchActivationQuantizationHolder
+
     def is_pytorch_layer_exportable(layer: Any) -> bool:
         """
         Check whether a torch Module is a valid exportable module or not.
@@ -31,12 +34,35 @@ if FOUND_TORCH:
         Returns:
             Check whether a PyTorch layer is a valid exportable layer or not.
         """
+        if not isinstance(layer, nn.Module):
+            Logger.error(f'Exportable layer must be a nn.Module layer, but layer {layer.name} is of type {type(layer)}') # pragma: no cover
+
         if isinstance(layer, PytorchQuantizationWrapper):
-            quantizers = list(layer.weights_quantizers.values())
-            quantizers.extend(layer.activation_quantizers)
-            if all([isinstance(q, BasePyTorchInferableQuantizer) for q in quantizers]):
-                return True
-        return False
+            valid_weights_quantizers = isinstance(layer.weights_quantizers, dict)
+            if not valid_weights_quantizers:
+                Logger.error(
+                    f'PytorchQuantizationWrapper must have a weights_quantizers but has a '
+                    f'{type(layer.weights_quantizers)} object') # pragma: no cover
+
+            if len(layer.weights_quantizers) == 0:
+                Logger.error(f'PytorchQuantizationWrapper must have at least one weight quantizer, but found {len(layer.weights_quantizers)} quantizers.'
+                             f'If layer is not quantized it should be a Keras layer.')
+
+            for _, weights_quantizer in layer.weights_quantizers.items():
+                if not isinstance(weights_quantizer, BasePyTorchInferableQuantizer):
+                    Logger.error(
+                        f'weights_quantizer must be a BasePyTorchInferableQuantizer object but has a '
+                        f'{type(weights_quantizer)} object')  # pragma: no cover
+
+        elif isinstance(layer, PytorchActivationQuantizationHolder):
+            if not isinstance(layer.activation_holder_quantizer, BasePyTorchInferableQuantizer):
+                Logger.error(
+                    f'activation quantizer in PytorchActivationQuantizationHolder'
+                    f' must be a BasePyTorchInferableQuantizer object but has a '
+                    f'{type(layer.activation_holder_quantizer)} object')  # pragma: no cover
+
+        return True
+
 else:
     def is_pytorch_layer_exportable(*args, **kwargs):  # pragma: no cover
         Logger.error('Installing torch is mandatory '

model_compression_toolkit/gptq/keras/gptq_training.py

@@ -26,8 +26,7 @@ from model_compression_toolkit.core.keras.back2framework.keras_model_builder imp
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
 from model_compression_toolkit.gptq.keras.quantizer.quantization_builder import quantization_builder
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.quantizers_infrastructure import KerasQuantizationWrapper
-from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.activation_quantization_holder import ActivationQuantizationHolder
+from mct_quantizers import KerasQuantizationWrapper, KerasActivationQuantizationHolder
 
 if version.parse(tf.__version__) < version.parse("2.6"):
     from tensorflow.python.keras.engine.base_layer import TensorFlowOpLayer
@@ -45,7 +44,6 @@ from model_compression_toolkit.core.common.framework_implementation import Frame
 import numpy as np
 import copy
 from model_compression_toolkit.core.keras.constants import BIAS, USE_BIAS
-from model_compression_toolkit import quantizers_infrastructure as qi
 
 
 class KerasGPTQTrainer(GPTQTrainer):
@@ -133,7 +131,7 @@ class KerasGPTQTrainer(GPTQTrainer):
 
     def gptq_wrapper(self,
                      n: common.BaseNode,
-                     layer: Layer) -> Union[qi.KerasQuantizationWrapper, Layer]:
+                     layer: Layer) -> Union[KerasQuantizationWrapper, Layer]:
         """
         A function which takes a computational graph node and a keras layer and perform the quantization wrapping.
 
@@ -145,22 +143,23 @@ class KerasGPTQTrainer(GPTQTrainer):
 
         """
         if self._is_gptq_weights_trainable(n):
-            weights_quantizers, _ = quantization_builder(n, self.gptq_config) # TODO: split quantizers building into two functions: for weights and activations
-            return qi.KerasQuantizationWrapper(layer,
-                                               weights_quantizers=weights_quantizers)
-        else:
-            return layer
-
-    def get_activation_quantizer_holder(self, n: common.BaseNode) -> Union[None, Callable]:
+            weights_quantizers, _ = quantization_builder(n,
+                                                         self.gptq_config) # TODO: split quantizers building into two functions: for weights and activations
+            if len(weights_quantizers) > 0:
+                return KerasQuantizationWrapper(layer,
+                                                   weights_quantizers=weights_quantizers)
+        return layer
+
+    def get_activation_quantizer_holder(self, n: common.BaseNode) -> 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.
 
         Returns:
-            A ActivationQuantizationHolder layer for the node activation quantization.
+            A KerasActivationQuantizationHolder layer for the node activation quantization.
         """
         _, activation_quantizers = quantization_builder(n, self.gptq_config) # TODO: split quantizers building into two functions: for weights and activations
 
@@ -168,10 +167,10 @@ class KerasGPTQTrainer(GPTQTrainer):
         # 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])
+            return KerasActivationQuantizationHolder(activation_quantizers[0])
 
         Logger.error(
-            f'ActivationQuantizationHolder supports a single quantizer but {len(activation_quantizers)} quantizers '
+            f'KerasActivationQuantizationHolder supports a single quantizer but {len(activation_quantizers)} quantizers '
             f'were found for node {n}')
 
 

model_compression_toolkit/gptq/keras/graph_info.py

@@ -21,8 +21,8 @@ from tensorflow.keras.models import Model
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.quantizers_infrastructure import KerasQuantizationWrapper
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
+from mct_quantizers import KerasQuantizationWrapper
+from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
 def get_gptq_trainable_parameters(fxp_model: Model,

model_compression_toolkit/gptq/keras/quantizer/base_keras_gptq_quantizer.py

@@ -19,15 +19,14 @@ from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import FOUND_TF
 from model_compression_toolkit.gptq.common.gptq_constants import WEIGHTS_QUANTIZATION_PARAMS
 
-from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig, \
-    TrainableQuantizerActivationConfig
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import BaseTrainableQuantizer
+from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig
+from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import BaseTrainableQuantizer
 
 if FOUND_TF:
     import tensorflow as tf
 
-    from model_compression_toolkit.quantizers_infrastructure import BaseKerasTrainableQuantizer, \
-        KerasQuantizationWrapper
+    from model_compression_toolkit.trainable_infrastructure import BaseKerasTrainableQuantizer
+    from mct_quantizers import KerasQuantizationWrapper
 
     class BaseKerasGPTQTrainableQuantizer(BaseKerasTrainableQuantizer):
         """

model_compression_toolkit/gptq/keras/quantizer/quantization_builder.py

@@ -21,14 +21,12 @@ from model_compression_toolkit.exporter.model_wrapper.keras.builder.node_to_quan
     get_inferable_quantizer_kwargs
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
 from model_compression_toolkit.gptq.keras.quantizer.base_keras_gptq_quantizer import BaseKerasGPTQTrainableQuantizer
-from model_compression_toolkit.quantizers_infrastructure import QuantizationTarget
-from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.get_quantizers import \
-    get_inferable_quantizer_class
-from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers.base_keras_inferable_quantizer import \
-    BaseKerasInferableQuantizer
-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.keras.quantizers import BaseKerasInferableQuantizer
+from model_compression_toolkit.trainable_infrastructure.common.get_quantizer_config import \
     get_trainable_quantizer_weights_config
-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/keras/quantizer/soft_rounding/soft_quantizer_reg.py

@@ -19,7 +19,7 @@ from keras import Model
 
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
-from model_compression_toolkit.quantizers_infrastructure import KerasQuantizationWrapper
+from mct_quantizers import KerasQuantizationWrapper
 
 
 class LinearTempDecay:

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

@@ -17,9 +17,9 @@ 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.core.common import max_power_of_two
 from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
+from mct_quantizers import QuantizationTarget
 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
@@ -27,11 +27,11 @@ from typing import Dict, Any
 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
-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: tf.Tensor,
@@ -66,7 +66,7 @@ def soft_rounding_symmetric_quantizer(input_tensor: tf.Tensor,
     return delta * 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.SoftQuantizer)
 class SymmetricSoftRoundingGPTQ(BaseKerasGPTQTrainableQuantizer):

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

@@ -17,20 +17,20 @@ 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.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
 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 \
+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_uniform_quantizer(input_tensor: tf.Tensor,
@@ -61,7 +61,7 @@ def soft_rounding_uniform_quantizer(input_tensor: tf.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(BaseKerasGPTQTrainableQuantizer):

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

@@ -19,18 +19,18 @@ import numpy as np
 import tensorflow as tf
 
 from model_compression_toolkit.gptq import RoundingType
-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
 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.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
-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 pertubation_symmetric_quantizer(input_tensor: tf.Tensor,
@@ -67,7 +67,7 @@ def pertubation_symmetric_quantizer(input_tensor: tf.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(BaseKerasGPTQTrainableQuantizer):

model_compression_toolkit/gptq/pytorch/gptq_training.py

@@ -32,9 +32,8 @@ from model_compression_toolkit.core.pytorch.utils import to_torch_tensor, set_mo
 from model_compression_toolkit.gptq.pytorch.graph_info import get_gptq_trainable_parameters, \
     get_weights_for_loss
 from model_compression_toolkit.gptq.pytorch.quantizer.quantization_builder import quantization_builder
-from model_compression_toolkit import quantizers_infrastructure as qi
 from model_compression_toolkit.gptq.pytorch.quantizer.regularization_factory import get_regularization
-from model_compression_toolkit.quantizers_infrastructure import PytorchQuantizationWrapper
+from mct_quantizers import PytorchQuantizationWrapper, PytorchActivationQuantizationHolder
 
 
 class PytorchGPTQTrainer(GPTQTrainer):
@@ -90,8 +89,8 @@ class PytorchGPTQTrainer(GPTQTrainer):
 
         self.reg_func = get_regularization(self.gptq_config, representative_data_gen)
 
-    def _is_gptq_applicable(self,
-                            node: BaseNode) -> bool:
+    def _is_gptq_weights_trainable(self,
+                                   node: BaseNode) -> bool:
         """
         A function for deciding if a layer should be fine-tuned during GPTQ.
         Args:
@@ -105,7 +104,9 @@ class PytorchGPTQTrainer(GPTQTrainer):
                          f"without a kernel isn't supported.")
         return node.is_weights_quantization_enabled()
 
-    def gptq_wrapper(self, n: BaseNode, layer: Module) -> Union[qi.PytorchQuantizationWrapper, Module]:
+    def gptq_wrapper(self,
+                     n: BaseNode,
+                     layer: Module) -> Union[PytorchQuantizationWrapper, Module]:
         """
         A function which takes a computational graph node and a pytorch layer and perform the quantization wrapping.
 
@@ -116,14 +117,32 @@ class PytorchGPTQTrainer(GPTQTrainer):
         Returns: Wrapped layer if the layer should be wrap, otherwise returns the layer as is.
         """
 
-        if self._is_gptq_applicable(n):
+        if self._is_gptq_weights_trainable(n):
             weights_quantizers, activation_quantizers = quantization_builder(n, self.gptq_config)
-            return qi.PytorchQuantizationWrapper(layer,
-                                                 weights_quantizers=weights_quantizers,
-                                                 activation_quantizers=activation_quantizers)
+            return PytorchQuantizationWrapper(layer,
+                                              weights_quantizers=weights_quantizers)
         else:
             return layer
 
+    def get_activation_quantizer_holder(self, n: BaseNode) -> Callable:
+        """
+        Retrieve a PytorchActivationQuantizationHolder layer to use for activation quantization of a node.
+        If the layer is not supposed to be wrapped with an activation quantizer - return None.
+        Args:
+            n: Node to attach a PytorchActivationQuantizationHolder to its output.
+        Returns:
+            A PytorchActivationQuantizationHolder module for the node's activation quantization.
+        """
+        _, activation_quantizers = quantization_builder(n, self.gptq_config)
+        # 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 we no activation
+        # quantization, which in this case has an empty list).
+        if len(activation_quantizers) == 1:
+            return PytorchActivationQuantizationHolder(activation_quantizers[0])
+        Logger.error(
+            f'PytorchActivationQuantizationHolder supports a single quantizer but {len(activation_quantizers)} quantizers '
+            f'were found for node {n}')
+
     def build_gptq_model(self):
         """
         Build the GPTQ model with QuantizationWrappers
@@ -134,7 +153,8 @@ class PytorchGPTQTrainer(GPTQTrainer):
                                                          append2output=self.compare_points,
                                                          fw_info=self.fw_info,
                                                          wrapper=self.gptq_wrapper,
-                                                         return_float_outputs=True).build_model()
+                                                         return_float_outputs=True,
+                                                         get_activation_quantizer_holder_fn=self.get_activation_quantizer_holder).build_model()
 
         return gptq_model, gptq_user_info
 

model_compression_toolkit/gptq/pytorch/graph_info.py

@@ -18,8 +18,9 @@ from typing import List
 from model_compression_toolkit.core.pytorch.constants import BIAS
 from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
-from model_compression_toolkit.quantizers_infrastructure import PytorchQuantizationWrapper
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
+from model_compression_toolkit.logger import Logger
+from mct_quantizers import PytorchQuantizationWrapper
+from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
 def get_gptq_trainable_parameters(fxp_model: nn.Module,
@@ -46,6 +47,8 @@ def get_gptq_trainable_parameters(fxp_model: nn.Module,
                                                                   fw_info=DEFAULT_PYTORCH_INFO)
 
             # collect trainable weights per quantizer
+            if kernel_attribute not in layer.weights_quantizers:
+                Logger.error(f'{kernel_attribute} was not found in weight quantizers of layer {layer.layer}')
             quantizer_trainable_weights = layer.weights_quantizers[kernel_attribute].get_trainable_variables(VariableGroup.WEIGHTS)
             quantizer_trainable_threshold = layer.weights_quantizers[kernel_attribute].get_trainable_variables(VariableGroup.QPARAMS)
             trainable_aux_weights.extend(quantizer_trainable_weights)

model_compression_toolkit/gptq/pytorch/quantization_facade.py

@@ -53,7 +53,8 @@ if FOUND_TORCH:
                                 optimizer: Optimizer = Adam([torch.Tensor([])], lr=LR_DEFAULT),
                                 optimizer_rest: Optimizer = Adam([torch.Tensor([])], lr=LR_REST_DEFAULT),
                                 loss: Callable = multiple_tensors_mse_loss,
-                                log_function: Callable = None) -> GradientPTQConfigV2:
+                                log_function: Callable = None,
+                                use_hessian_based_weights: bool = True) -> GradientPTQConfigV2:
         """
         Create a GradientPTQConfigV2 instance for Pytorch models.
 
@@ -63,6 +64,7 @@ if FOUND_TORCH:
             optimizer_rest (Optimizer): Pytorch optimizer to use for fine-tuning of the bias variable.
             loss (Callable): loss to use during fine-tuning. should accept 4 lists of tensors. 1st list of quantized tensors, the 2nd list is the float tensors, the 3rd is a list of quantized weights and the 4th is a list of float weights.
             log_function (Callable): Function to log information about the gptq process.
+            use_hessian_based_weights (bool): Whether to use Hessian-based weights for weighted average loss.
 
         returns:
             a GradientPTQConfigV2 object to use when fine-tuning the quantized model using gptq.
@@ -84,7 +86,7 @@ if FOUND_TORCH:
         """
         bias_optimizer = torch.optim.SGD([torch.Tensor([])], lr=LR_BIAS_DEFAULT, momentum=GPTQ_MOMENTUM)
         return GradientPTQConfigV2(n_epochs, optimizer, optimizer_rest=optimizer_rest, loss=loss,
-                                   log_function=log_function, train_bias=True, optimizer_bias=bias_optimizer)
+                                   log_function=log_function, train_bias=True, optimizer_bias=bias_optimizer, use_hessian_based_weights=use_hessian_based_weights)
 
 
     def pytorch_gradient_post_training_quantization_experimental(model: Module,

model_compression_toolkit/gptq/pytorch/quantizer/base_pytorch_gptq_quantizer.py

@@ -19,16 +19,16 @@ 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, \
+from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig, \
     TrainableQuantizerActivationConfig
-from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import \
+from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import \
     BaseTrainableQuantizer
-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:
     from torch import Tensor
-    from model_compression_toolkit.quantizers_infrastructure import PytorchQuantizationWrapper
+    from mct_quantizers import PytorchQuantizationWrapper
 
     class BasePytorchGPTQTrainableQuantizer(BasePytorchTrainableQuantizer):
         """