mct-nightly 2.2.0.20241011.452__py3-none-any.whl → 2.2.0.20241017.455__py3-none-any.whl

{mct_nightly-2.2.0.20241011.452.dist-info → mct_nightly-2.2.0.20241017.455.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.2.0.20241011.452
+Version: 2.2.0.20241017.455
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct_nightly-2.2.0.20241011.452.dist-info → mct_nightly-2.2.0.20241017.455.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=9L9kR-Q7DX8jQeSpnCtsLriXOdDH-fNxLDaLy1zcpPY,1573
+model_compression_toolkit/__init__.py,sha256=Pw5laVBBrwaRaU8RB7HsipMx-1faRX33B0Lfa6YX2kA,1573
 model_compression_toolkit/constants.py,sha256=i4wYheBkIdQmsQA-axIpcT3YiSO1USNc-jaNiNE8w6E,3920
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
@@ -402,15 +402,15 @@ model_compression_toolkit/qat/common/qat_config.py,sha256=xtfVSoyELGXynHNrw86dB9
 model_compression_toolkit/qat/keras/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/qat/keras/quantization_facade.py,sha256=LNM2HW4cNei3tUhwLdNtsWrox_uSAhaswFxWiMEIrPM,17278
 model_compression_toolkit/qat/keras/quantizer/__init__.py,sha256=zmYyCa25_KLCSUCGUDRslh3RCIjcRMxc_oXa54Aui-4,996
-model_compression_toolkit/qat/keras/quantizer/base_keras_qat_quantizer.py,sha256=hoY3AETaLSRP7YfecZ32tyUUj-X_DHRWkV8nALYeRlY,2202
+model_compression_toolkit/qat/keras/quantizer/base_keras_qat_weight_quantizer.py,sha256=EbIt4lMlh6cU4awFLMBp0IlZ2zUUp-WtnlW5Wn19FDM,1793
 model_compression_toolkit/qat/keras/quantizer/quant_utils.py,sha256=cBULOgWUodcBO1lHevZggdTevuDYI6tQceV86U2x6DA,2543
-model_compression_toolkit/qat/keras/quantizer/quantization_builder.py,sha256=HD0JIOiqnrpqj5qk6RyzuCsSGZsDUVohdCYSePmJBNQ,5872
+model_compression_toolkit/qat/keras/quantizer/quantization_builder.py,sha256=hGizGBbOGZpD-w3wg-LlehUYJDWLk91VUdfVwwG2Z78,5882
 model_compression_toolkit/qat/keras/quantizer/lsq/__init__.py,sha256=lNJ29DYxaLUPDstRDA1PGI5r9Fulq_hvrZMlhst1Z5g,697
-model_compression_toolkit/qat/keras/quantizer/lsq/symmetric_lsq.py,sha256=MwHo4qUYTm-cZZ9f4bEDU2fcdO1VdLXcrp8MKhJ051k,12043
-model_compression_toolkit/qat/keras/quantizer/lsq/uniform_lsq.py,sha256=lGMJF_8jgHV2Rp97aMIqt7B7Gn7JsEOVbBW55K9tvuI,11244
+model_compression_toolkit/qat/keras/quantizer/lsq/symmetric_lsq.py,sha256=dtkS0mpjvJntAxpOi-BJx-pCeBF2ReKKeH7y2uwzpH0,6756
+model_compression_toolkit/qat/keras/quantizer/lsq/uniform_lsq.py,sha256=vGUs9b0IHTydCA5tN7iekuhf1LHNgIrSF5sXMD1WsSI,6476
 model_compression_toolkit/qat/keras/quantizer/ste_rounding/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/qat/keras/quantizer/ste_rounding/symmetric_ste.py,sha256=fPAC49mBlB5ViaQT_xHUTC8EvH84OsBX3WAPusqYcM8,13538
-model_compression_toolkit/qat/keras/quantizer/ste_rounding/uniform_ste.py,sha256=6YS0v1qCq5dRqtLKHc2gHaKJWfql84TxtZ7pypaZock,10810
+model_compression_toolkit/qat/keras/quantizer/ste_rounding/symmetric_ste.py,sha256=pFkrjtlavCniswcO3-Djlh6a_Hz1rrcEa7Z5wTGVRCU,8270
+model_compression_toolkit/qat/keras/quantizer/ste_rounding/uniform_ste.py,sha256=ZdZwMwLa1Ws2eo3DiQYYTvPS1JfiswZL1xlQPtRnIgE,7067
 model_compression_toolkit/qat/pytorch/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/qat/pytorch/quantization_facade.py,sha256=NnFy2E_7SR2m8vfh8Q8VrXOXhe7rMScgXnYBtDpsqVs,13456
 model_compression_toolkit/qat/pytorch/quantizer/__init__.py,sha256=xYa4C8pr9cG1f3mQQcBXO_u3IdJN-zl7leZxuXDs86w,1003
@@ -500,11 +500,19 @@ model_compression_toolkit/trainable_infrastructure/common/quant_utils.py,sha256=
 model_compression_toolkit/trainable_infrastructure/common/trainable_quantizer_config.py,sha256=My5Wz34jPOyh8z33OTpKnOobRB0cpO_Qgmtsd5lizHo,4791
 model_compression_toolkit/trainable_infrastructure/common/training_method.py,sha256=LUoeJkloowhZKuHTiOfzjmSUn2G-4of11-rbnL-h0P4,1194
 model_compression_toolkit/trainable_infrastructure/keras/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
-model_compression_toolkit/trainable_infrastructure/keras/base_keras_quantizer.py,sha256=tHEI9vkLjBzdeCD7eTgAHuUubmnq8GbWSF7Coun8zzE,4116
+model_compression_toolkit/trainable_infrastructure/keras/base_keras_quantizer.py,sha256=LBc26z8pkpbcdKMTxpNBg5IyChLreHQ1lRgCVjNE37o,4202
 model_compression_toolkit/trainable_infrastructure/keras/config_serialization.py,sha256=txdWXdZoHazg-3MDPb9P-oXRM92LRn2G_8woEplwKaI,4360
 model_compression_toolkit/trainable_infrastructure/keras/load_model.py,sha256=DJHibcLo-UCuHV6UPLeVd7dKmPfkGXEiLqCCqvQrISM,3769
 model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py,sha256=eVB5FSE3OmTLrhfLUcP2knwN1z2_unQLM-xFEGwdafA,5587
-model_compression_toolkit/trainable_infrastructure/keras/quantizer_utils.py,sha256=MVwXNymmFRB2NXIBx4e2mdJ1RfoHxRPYRgjb1MQP5kY,1797
+model_compression_toolkit/trainable_infrastructure/keras/quantizer_utils.py,sha256=r3CaPd4pyM1GDXU2--9NT3wwvl9H6y3QUrVT9spx5es,4189
+model_compression_toolkit/trainable_infrastructure/keras/activation_quantizers/__init__.py,sha256=QPBRTl_9ZXF-Yk5srotlKVOmxKTXMm5xf2-9IjIrBAI,1055
+model_compression_toolkit/trainable_infrastructure/keras/activation_quantizers/base_activation_quantizer.py,sha256=VvwsrOVZgWed82P9rtu_UDDD99MnZSppPsjrCtxk2AY,964
+model_compression_toolkit/trainable_infrastructure/keras/activation_quantizers/lsq/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
+model_compression_toolkit/trainable_infrastructure/keras/activation_quantizers/lsq/symmetric_lsq.py,sha256=iBO2hAwJ1HUQzsis_kEGE-BPB3hOW7IF7p2uFLWg09A,6259
+model_compression_toolkit/trainable_infrastructure/keras/activation_quantizers/lsq/uniform_lsq.py,sha256=2BOQXymCZUSLdxDbaS8Blr2FB-NxQV01punWNjMGiNc,5765
+model_compression_toolkit/trainable_infrastructure/keras/activation_quantizers/ste/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
+model_compression_toolkit/trainable_infrastructure/keras/activation_quantizers/ste/symmetric_ste.py,sha256=THY5eZ_69D1yzkXLhLg84ON_deNUAD_qMJ6A5C5znDM,7359
+model_compression_toolkit/trainable_infrastructure/keras/activation_quantizers/ste/uniform_ste.py,sha256=XEypYorBnSBLj6sh1pHCNaSjeCToYVlERWIHxUoXvuc,5733
 model_compression_toolkit/trainable_infrastructure/pytorch/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/trainable_infrastructure/pytorch/annealing_schedulers.py,sha256=IdUBpZUcOXHLPp2OhwbO_Kytee3OTVuy2032N-tm694,1686
 model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py,sha256=lWc5EG3ptrP85n69EHGKFkIadnrKEBMKnB5YXQ5AmXo,2745
@@ -551,8 +559,8 @@ tests_pytest/pytorch/gptq/test_annealing_cfg.py,sha256=hGC7L6mp3N1ygcJ3OctgS_Fz2
 tests_pytest/pytorch/gptq/test_gradual_act_quantization.py,sha256=tI01aFIUaiCILL5Qn--p1E_rLBUelxLdSY3k52lwcx0,4594
 tests_pytest/pytorch/trainable_infrastructure/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
 tests_pytest/pytorch/trainable_infrastructure/test_linear_annealing.py,sha256=eNOpSp0GoLxtEdiRypBp8jaujXfdNxBwKh5Rd-P7WLs,1786
-mct_nightly-2.2.0.20241011.452.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.2.0.20241011.452.dist-info/METADATA,sha256=gUF-VCq67b82gtIGg40WsJpMhQhOsE5JKTUAzy6KbxI,20830
-mct_nightly-2.2.0.20241011.452.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-mct_nightly-2.2.0.20241011.452.dist-info/top_level.txt,sha256=csdfSXhtRnpWYRzjZ-dRLIhOmM2TEdVXUxG05A5fgb8,39
-mct_nightly-2.2.0.20241011.452.dist-info/RECORD,,
+mct_nightly-2.2.0.20241017.455.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.2.0.20241017.455.dist-info/METADATA,sha256=OkKILNYbh3SMCjpEVFD8rXhbjfcGeksV0H7PG-ug_d4,20830
+mct_nightly-2.2.0.20241017.455.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+mct_nightly-2.2.0.20241017.455.dist-info/top_level.txt,sha256=csdfSXhtRnpWYRzjZ-dRLIhOmM2TEdVXUxG05A5fgb8,39
+mct_nightly-2.2.0.20241017.455.dist-info/RECORD,,

model_compression_toolkit/__init__.py

@@ -27,4 +27,4 @@ from model_compression_toolkit import data_generation
 from model_compression_toolkit import pruning
 from model_compression_toolkit.trainable_infrastructure.keras.load_model import keras_load_quantized_model
 
-__version__ = "2.2.0.20241011.000452"
+__version__ = "2.2.0.20241017.000455"

model_compression_toolkit/qat/keras/quantizer/{base_keras_qat_quantizer.py → base_keras_qat_weight_quantizer.py}

@@ -22,24 +22,14 @@ from model_compression_toolkit.trainable_infrastructure import TrainableQuantize
 
 if FOUND_TF:
 
-    class BaseKerasQATTrainableQuantizer(BaseKerasTrainableQuantizer):
+    class BaseKerasQATWeightTrainableQuantizer(BaseKerasTrainableQuantizer):
         """
         A base class for trainable Keras quantizer for QAT.
         """
-
-        def __init__(self,
-                     quantization_config: Union[TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig]):
-            """
-            Initializes BaseKerasQATTrainableQuantizer object.
-
-            Args:
-                quantization_config: quantizer config class contains all the information about a quantizer configuration.
-            """
-
-            super().__init__(quantization_config)
+        pass
 
 else:  # pragma: no cover
-    class BaseKerasQATTrainableQuantizer(BaseKerasTrainableQuantizer):
+    class BaseKerasQATWeightTrainableQuantizer(BaseKerasTrainableQuantizer):
         def __init__(self,
                      quantization_config: Union[TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig]):
 

model_compression_toolkit/qat/keras/quantizer/lsq/symmetric_lsq.py

@@ -28,47 +28,18 @@ from mct_quantizers import QuantizationTarget, mark_quantizer
 from model_compression_toolkit.qat.common import THRESHOLD_TENSOR
 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.trainable_infrastructure import TrainableQuantizerWeightsConfig, \
-    TrainableQuantizerActivationConfig
-from mct_quantizers.keras.quantizers import WeightsPOTInferableQuantizer, WeightsSymmetricInferableQuantizer, \
-    ActivationPOTInferableQuantizer, ActivationSymmetricInferableQuantizer
+from model_compression_toolkit.qat.keras.quantizer.base_keras_qat_weight_quantizer import BaseKerasQATWeightTrainableQuantizer
+from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig
+from mct_quantizers.keras.quantizers import WeightsPOTInferableQuantizer, WeightsSymmetricInferableQuantizer, ActivationPOTInferableQuantizer, ActivationSymmetricInferableQuantizer
 from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 from model_compression_toolkit.qat.keras.quantizer.quant_utils import ste_round, grad_scale
-
-
-def symmetric_lsq_quantizer(x: tf.Tensor,
-                            thresholds: tf.Tensor,
-                            num_bits: int,
-                            sign: bool,
-                            min_int: int,
-                            max_int:int,
-                            scale_factor: float) -> tf.Tensor:
-    """
-    Symmetric quantizer according to LSQ algorithm: https://arxiv.org/pdf/1902.08153.pdf
-    Args:
-        x: input to quantize
-        thresholds: thresholds of quantization levels
-        num_bits: number of bits for quantization
-        sign: whether x is signed or not
-        min_int: min clipping integer value
-        max_int: max clipping integer value
-        scale_factor: grad scale of LSQ algorithm
-    Returns:
-        A quantized tensor
-    """
-    delta = thresholds / (2 ** (num_bits - int(sign)))
-    delta_scaled = grad_scale(delta, scale_factor)
-    rounded = ste_round(x / delta_scaled)
-    clipped = tf.math.minimum(tf.math.maximum(rounded, min_int), max_int)
-    quantized = delta_scaled * clipped
-    return quantized
+from model_compression_toolkit.trainable_infrastructure.keras.quantizer_utils import symmetric_lsq_quantizer
 
 
 @mark_quantizer(quantization_target=QuantizationTarget.Weights,
                 quantization_method=[QuantizationMethod.POWER_OF_TWO, QuantizationMethod.SYMMETRIC],
                 identifier=TrainingMethod.LSQ)
-class LSQWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
+class LSQWeightQATQuantizer(BaseKerasQATWeightTrainableQuantizer):
     """
     Trainable constrained quantizer to quantize layer's weights.
     """
@@ -159,95 +130,3 @@ class LSQWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
                                                       input_rank=len(self.threshold_shape))
 
 
-@mark_quantizer(quantization_target=QuantizationTarget.Activation,
-                quantization_method=[QuantizationMethod.POWER_OF_TWO, QuantizationMethod.SYMMETRIC],
-                identifier=TrainingMethod.LSQ)
-class LSQActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
-    """
-    Trainable constrained quantizer to quantize layer activations.
-    """
-
-    def __init__(self, quantization_config: TrainableQuantizerActivationConfig):
-        """
-        Initialize a LSQActivationQATQuantizer object with parameters to use
-        for the quantization.
-
-        Args:
-            quantization_config: trainable quantizer config class
-        """
-        super().__init__(quantization_config)
-        self.power_of_two = quantization_config.activation_quantization_method == QuantizationMethod.POWER_OF_TWO
-        self.threshold_values = float(quantization_config.activation_quantization_params[C.THRESHOLD])
-        self.threshold_shape = np.asarray(self.threshold_values).shape
-        self.sign = quantization_config.activation_quantization_params[SIGNED]
-        self.num_bits = quantization_config.activation_n_bits
-        n_pos_bits = self.num_bits - int(self.sign)
-        self.min_int = -int(self.sign) * (2 ** n_pos_bits)
-        self.max_int = (2 ** n_pos_bits) - 1
-        if self.power_of_two:
-            self.threshold_values = np.power(2.0, np.ceil(np.log2(np.maximum(self.threshold_values, C.MIN_THRESHOLD))))
-
-
-    def initialize_quantization(self,
-                                tensor_shape: TensorShape,
-                                name: str,
-                                layer: KerasTrainableQuantizationWrapper):
-        """
-        Add quantizer parameters to the quantizer parameters dictionary
-
-        Args:
-            tensor_shape: tensor shape of the quantized tensor.
-            name: Tensor name.
-            layer: Layer to quantize.
-        """
-        ptq_threshold_tensor = layer.add_weight(
-            name + THRESHOLD_TENSOR,
-            shape=(),
-            initializer=tf.keras.initializers.Constant(1.0),
-            trainable=True)
-        ptq_threshold_tensor.assign(self.threshold_values)
-
-        # save the quantizer added parameters for later calculations
-        self.add_quantizer_variable(THRESHOLD_TENSOR, ptq_threshold_tensor, VariableGroup.QPARAMS)
-
-    def __call__(self,
-                 inputs: tf.Tensor,
-                 training: bool):
-        """
-        Quantize a tensor.
-        Args:
-            inputs: Input tensor to quantize.
-            training: Whether the graph is in training mode.
-
-        Returns:
-            The quantized tensor.
-        """
-
-        thresholds = self.get_quantizer_variable(THRESHOLD_TENSOR)
-        n_channels = inputs.shape[-1]
-        scale_factor = 1.0 / np.sqrt(self.max_int * n_channels)
-        q_tensor = symmetric_lsq_quantizer(inputs, thresholds, self.num_bits, self.sign, self.min_int, self.max_int, scale_factor)
-        return q_tensor
-
-    def convert2inferable(self) -> Union[ActivationPOTInferableQuantizer, ActivationSymmetricInferableQuantizer]:
-        """
-        Convert quantizer to inferable quantizer.
-
-        Returns:
-            BaseKerasInferableQuantizer object.
-        """
-
-        if self.power_of_two:
-            thresholds = 2 ** np.ceil(np.log2(self.get_quantizer_variable(THRESHOLD_TENSOR).numpy()))
-            return ActivationPOTInferableQuantizer(num_bits=self.num_bits,
-                                                   # In activation quantization is per-tensor only - thus we pass
-                                                   # the threshold as a list with a len of 1
-                                                   threshold=[thresholds],
-                                                   signed=self.sign)
-        else:
-            thresholds = self.get_quantizer_variable(THRESHOLD_TENSOR).numpy()
-            return ActivationSymmetricInferableQuantizer(num_bits=self.num_bits,
-                                                         # In activation quantization is per-tensor only - thus we
-                                                         # pass the threshold as a list with a len of 1
-                                                         threshold=[thresholds],
-                                                         signed=self.sign)

model_compression_toolkit/qat/keras/quantizer/lsq/uniform_lsq.py

@@ -16,6 +16,8 @@ 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.qat.keras.quantizer.base_keras_qat_weight_quantizer import \
+    BaseKerasQATWeightTrainableQuantizer
 from model_compression_toolkit.trainable_infrastructure.common.constants import FQ_MIN, FQ_MAX
 from model_compression_toolkit.trainable_infrastructure import KerasTrainableQuantizationWrapper
 from model_compression_toolkit.trainable_infrastructure import TrainingMethod
@@ -26,47 +28,18 @@ from mct_quantizers.keras.quantizers import \
 
 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.trainable_infrastructure import TrainableQuantizerWeightsConfig, \
     TrainableQuantizerActivationConfig
 from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import fix_range_to_include_zero
 from model_compression_toolkit.qat.keras.quantizer.quant_utils import ste_round, grad_scale, adjust_range_to_include_zero
-
-
-def uniform_lsq_quantizer(x: tf.Tensor,
-                          min_range: tf.Tensor,
-                          max_range: tf.Tensor,
-                          num_bits: int,
-                          min_int: int,
-                          max_int:int,
-                          scale_factor: float) -> tf.Tensor:
-    """
-    Uniform quantizer according to LSQ algorithm: https://arxiv.org/pdf/1902.08153.pdf
-    Args:
-        x: input to quantize
-        min_range: min range of quantization values
-        max_range: min range of quantization values
-        num_bits: number of bits for quantization
-        min_int: min clipping integer value
-        max_int: max clipping integer value
-        scale_factor: grad scale of LSQ algorithm
-    Returns:
-        A quantized tensor
-    """
-    min_range, max_range = adjust_range_to_include_zero(min_range, max_range, num_bits)
-    delta = (max_range - min_range) / (2 ** num_bits - 1)
-    delta_scaled = grad_scale(delta, scale_factor)
-    rounded = ste_round((x-min_range) / delta_scaled)
-    clipped = tf.math.minimum(tf.math.maximum(rounded, min_int), max_int)
-    quantized = delta_scaled * clipped + min_range
-    return quantized
+from model_compression_toolkit.trainable_infrastructure.keras.quantizer_utils import uniform_lsq_quantizer
 
 
 @mark_quantizer(quantization_target=QuantizationTarget.Weights,
                 quantization_method=[QuantizationMethod.UNIFORM],
                 identifier=TrainingMethod.LSQ)
-class LSQUniformWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
+class LSQUniformWeightQATQuantizer(BaseKerasQATWeightTrainableQuantizer):
     """
     Trainable constrained quantizer to quantize layer's weights.
     """
@@ -158,93 +131,3 @@ class LSQUniformWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
                                                 channel_axis=self.channel_axis,
                                                 input_rank=len(self.min_max_shape))
 
-
-@mark_quantizer(quantization_target=QuantizationTarget.Activation,
-                quantization_method=[QuantizationMethod.UNIFORM],
-                identifier=TrainingMethod.LSQ)
-class LSQUniformActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
-    """
-    Trainable constrained quantizer to quantize layer activations.
-    """
-
-    def __init__(self, quantization_config: TrainableQuantizerActivationConfig):
-        """
-        Initialize a LSQUniformActivationQATQuantizer object with parameters to use
-        for the quantization.
-
-        Args:
-            quantization_config: trainable quantizer config class
-        """
-        super().__init__(quantization_config)
-
-        self.num_bits = quantization_config.activation_n_bits
-        self.min_range = np.array(quantization_config.activation_quantization_params[C.RANGE_MIN])
-        self.max_range = np.array(quantization_config.activation_quantization_params[C.RANGE_MAX])
-        self.min_int = 0
-        self.max_int = 2**self.num_bits - 1
-
-    def initialize_quantization(self,
-                                tensor_shape: TensorShape,
-                                name: str,
-                                layer: KerasTrainableQuantizationWrapper):
-        """
-        Add quantizer parameters to the quantizer parameters dictionary
-
-        Args:
-            tensor_shape: tensor shape of the quantized tensor.
-            name: Tensor name.
-            layer: Layer to quantize.
-        """
-        fq_min = layer.add_weight(
-            name + FQ_MIN,
-            shape=(),
-            initializer=tf.keras.initializers.Constant(-1.0),
-            trainable=True)
-        fq_min.assign(self.min_range)
-
-        fq_max = layer.add_weight(
-            name + FQ_MAX,
-            shape=(),
-            initializer=tf.keras.initializers.Constant(1.0),
-            trainable=True)
-        fq_max.assign(self.max_range)
-
-        # save the quantizer added parameters for later calculations
-        self.add_quantizer_variable(FQ_MIN, fq_min, VariableGroup.QPARAMS)
-        self.add_quantizer_variable(FQ_MAX, fq_max, VariableGroup.QPARAMS)
-
-    def __call__(self,
-                 inputs: tf.Tensor,
-                 training: bool):
-        """
-        Quantize a tensor.
-        Args:
-            inputs: Input tensor to quantize.
-            training: Whether the graph is in training mode.
-
-        Returns:
-            The quantized tensor.
-        """
-
-        min_range = self.get_quantizer_variable(FQ_MIN)
-        max_range = self.get_quantizer_variable(FQ_MAX)
-        n_channels = inputs.shape[-1]
-        scale_factor = 1.0 / np.sqrt(self.max_int * n_channels)
-        q_tensor = uniform_lsq_quantizer(inputs, min_range, max_range, self.num_bits, self.min_int, self.max_int, scale_factor)
-        return q_tensor
-
-    def convert2inferable(self) -> BaseKerasInferableQuantizer:
-        """
-        Convert quantizer to inferable quantizer.
-
-        Returns:
-            BaseKerasInferableQuantizer object.
-        """
-        min_range, max_range = fix_range_to_include_zero(self.get_quantizer_variable(FQ_MIN).numpy(),
-                                                         self.get_quantizer_variable(FQ_MAX).numpy(),
-                                                         self.num_bits)
-        return ActivationUniformInferableQuantizer(num_bits=self.num_bits,
-                                                   # In activation quantization is per-tensor only - thus we pass
-                                                   # the min/max as lists with a len of 1
-                                                   min_range=[min_range],
-                                                   max_range=[max_range])

model_compression_toolkit/qat/keras/quantizer/quantization_builder.py

@@ -15,17 +15,18 @@
 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
-from model_compression_toolkit.qat.keras.quantizer.base_keras_qat_quantizer import BaseKerasQATTrainableQuantizer
 from mct_quantizers import QuantizationTarget, KerasActivationQuantizationHolder
+from model_compression_toolkit.qat.keras.quantizer.base_keras_qat_weight_quantizer import \
+    BaseKerasQATWeightTrainableQuantizer
 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.trainable_infrastructure.common.get_quantizers import \
     get_trainable_quantizer_class
+from model_compression_toolkit.trainable_infrastructure.keras.activation_quantizers import \
+    BaseKerasActivationTrainableQuantizer
 
 
 def get_activation_quantizer_holder(n: common.BaseNode,
@@ -55,7 +56,7 @@ def get_activation_quantizer_holder(n: common.BaseNode,
 def quantization_builder(n: common.BaseNode,
                          qat_config: QATConfig,
                          kernel_attr: str = None,
-                         ) -> Tuple[Dict[str, BaseKerasQATTrainableQuantizer], List[BaseKerasQATTrainableQuantizer]]:
+                         ) -> Tuple[Dict[str, BaseKerasQATWeightTrainableQuantizer], List[BaseKerasActivationTrainableQuantizer]]:
     """
     Build quantizers for a node according to its quantization configuration.
 
@@ -82,7 +83,7 @@ def quantization_builder(n: common.BaseNode,
         quantizer_class = get_trainable_quantizer_class(QuantizationTarget.Weights,
                                                         qat_config.weight_training_method,
                                                         quant_method,
-                                                        BaseKerasQATTrainableQuantizer)
+                                                        BaseKerasQATWeightTrainableQuantizer)
 
         weight_quantizers.update({kernel_attr: quantizer_class(get_trainable_quantizer_weights_config(n,
                                                                                                       attr_name=kernel_attr,
@@ -98,7 +99,7 @@ def quantization_builder(n: common.BaseNode,
         quantizer_class = get_trainable_quantizer_class(QuantizationTarget.Activation,
                                                         qat_config.activation_training_method,
                                                         quant_method,
-                                                        BaseKerasQATTrainableQuantizer)
+                                                        BaseKerasActivationTrainableQuantizer)
 
         activation_quantizers = [quantizer_class(get_trainable_quantizer_activation_config(n, aq_cand),
                                                  **qat_config.activation_quantizer_params_override)] * len(output_shapes)

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

@@ -18,7 +18,6 @@ from typing import Union
 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.trainable_infrastructure.common.constants import FQ_MIN, FQ_MAX
 
 from model_compression_toolkit.trainable_infrastructure import TrainingMethod
@@ -29,18 +28,16 @@ from mct_quantizers import QuantizationTarget, mark_quantizer
 from model_compression_toolkit.qat.common import THRESHOLD_TENSOR
 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.trainable_infrastructure import TrainableQuantizerWeightsConfig, \
-    TrainableQuantizerActivationConfig
-from mct_quantizers.keras.quantizers import WeightsPOTInferableQuantizer, WeightsSymmetricInferableQuantizer, \
-    ActivationPOTInferableQuantizer, ActivationSymmetricInferableQuantizer
+from model_compression_toolkit.qat.keras.quantizer.base_keras_qat_weight_quantizer import BaseKerasQATWeightTrainableQuantizer
+from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig
+from mct_quantizers.keras.quantizers import WeightsPOTInferableQuantizer, WeightsSymmetricInferableQuantizer
 from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
 @mark_quantizer(quantization_target=QuantizationTarget.Weights,
                 quantization_method=[QuantizationMethod.POWER_OF_TWO, QuantizationMethod.SYMMETRIC],
                 identifier=TrainingMethod.STE)
-class STEWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
+class STEWeightQATQuantizer(BaseKerasQATWeightTrainableQuantizer):
     """
     Trainable constrained quantizer to quantize a layer inputs.
     """
@@ -171,115 +168,3 @@ class STEWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
                                                       input_rank=len(self.threshold_shape))
 
 
-@mark_quantizer(quantization_target=QuantizationTarget.Activation,
-                quantization_method=[QuantizationMethod.POWER_OF_TWO, QuantizationMethod.SYMMETRIC],
-                identifier=TrainingMethod.STE)
-class STEActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
-    """
-    Trainable constrained quantizer to quantize a layer outputs.
-    """
-
-    def __init__(self, quantization_config: TrainableQuantizerActivationConfig):
-        """
-        Initialize a STEActivationQATQuantizer object with parameters to use
-        for the quantization.
-
-        Args:
-            quantization_config: trainable quantizer config class
-        """
-        super().__init__(quantization_config)
-        self.power_of_two = quantization_config.activation_quantization_method == QuantizationMethod.POWER_OF_TWO
-        self.threshold_values = quantization_config.activation_quantization_params[C.THRESHOLD]
-        self.threshold_shape = np.asarray(self.threshold_values).shape
-        self.np_threshold_values = float(self.threshold_values)
-        self.signed = quantization_config.activation_quantization_params[SIGNED]
-        if self.power_of_two:
-            self.np_threshold_values = np.power(2.0,
-                                                np.ceil(np.log2(np.maximum(self.np_threshold_values, C.MIN_THRESHOLD))))
-        self.num_bits = quantization_config.activation_n_bits
-        delta = self.np_threshold_values / np.power(2.0, self.num_bits - int(self.signed))
-        min_int = -int(self.signed) * (2 ** (self.num_bits - int(self.signed)))
-        max_int = (2 ** (self.num_bits - int(self.signed))) - 1
-        self.min = delta * min_int
-        self.max = delta * max_int
-
-    def initialize_quantization(self,
-                                tensor_shape: TensorShape,
-                                name: str,
-                                layer: KerasTrainableQuantizationWrapper):
-        """
-        Add quantizer parameters to the quantizer parameters dictionary
-
-        Args:
-            tensor_shape: tensor shape of the quantized tensor.
-            name: Tensor name.
-            layer: Layer to quantize.
-        """
-        ptq_threshold_tensor = layer.add_weight(
-            name + THRESHOLD_TENSOR,
-            shape=(),
-            initializer=tf.keras.initializers.Constant(1.0),
-            trainable=False)
-        ptq_threshold_tensor.assign(self.np_threshold_values)
-
-        fq_min = layer.add_weight(
-            name + FQ_MIN,
-            shape=(),
-            initializer=tf.keras.initializers.Constant(-1.0),
-            trainable=False)
-        fq_min.assign(self.min)
-
-        fq_max = layer.add_weight(
-            name + FQ_MAX,
-            shape=(),
-            initializer=tf.keras.initializers.Constant(1.0),
-            trainable=False)
-        fq_max.assign(self.max)
-
-        # save the quantizer added parameters for later calculations
-        self.add_quantizer_variable(THRESHOLD_TENSOR, ptq_threshold_tensor, VariableGroup.QPARAMS)
-        self.add_quantizer_variable(FQ_MIN, fq_min, VariableGroup.QPARAMS)
-        self.add_quantizer_variable(FQ_MAX, fq_max, VariableGroup.QPARAMS)
-
-
-    def __call__(self,
-                 inputs: tf.Tensor,
-                 training: bool):
-        """
-        Quantize a tensor.
-        Args:
-            inputs: Input tensor to quantize.
-            training: Whether the graph is in training mode.
-
-        Returns:
-            The quantized tensor.
-        """
-
-        _min = self.get_quantizer_variable(FQ_MIN)
-        _max = self.get_quantizer_variable(FQ_MAX)
-        q_tensor = tf.quantization.fake_quant_with_min_max_vars(inputs, _min, _max,
-                                                                num_bits=self.num_bits)
-
-        return q_tensor
-
-    def convert2inferable(self) -> Union[ActivationPOTInferableQuantizer, ActivationSymmetricInferableQuantizer]:
-        """
-        Convert quantizer to inferable quantizer.
-
-        Returns:
-            BaseKerasInferableQuantizer object.
-        """
-
-        if self.power_of_two:
-            pot_threshold = 2 ** np.ceil(np.log2(self.get_quantizer_variable(THRESHOLD_TENSOR)))
-            return ActivationPOTInferableQuantizer(num_bits=self.num_bits,
-                                                      # In activation quantization is per-tensor only - thus we pass
-                                                      # the threshold as a list with a len of 1
-                                                      threshold=[pot_threshold],
-                                                      signed=self.signed)
-        else:
-            return ActivationSymmetricInferableQuantizer(num_bits=self.num_bits,
-                                                         # In activation quantization is per-tensor only - thus we
-                                                         # pass the threshold as a list with a len of 1
-                                                         threshold=[self.get_quantizer_variable(THRESHOLD_TENSOR).numpy()],
-                                                         signed=self.signed)