mct-nightly 2.2.0.20240902.511__py3-none-any.whl → 2.2.0.20240904.449__py3-none-any.whl

model_compression_toolkit/qat/pytorch/quantizer/quantization_builder.py

@@ -17,16 +17,17 @@ 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.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.qat.pytorch.quantizer.base_pytorch_qat_weight_quantizer import BasePytorchQATWeightTrainableQuantizer
 from model_compression_toolkit.trainable_infrastructure.common.get_quantizers import \
     get_trainable_quantizer_class
+from model_compression_toolkit.trainable_infrastructure.pytorch.activation_quantizers.base_activation_quantizer import \
+    BasePytorchActivationTrainableQuantizer
+
 
 def get_activation_quantizer_holder(n: common.BaseNode,
                                     qat_config: QATConfig) -> Callable:
@@ -55,8 +56,8 @@ def get_activation_quantizer_holder(n: common.BaseNode,
 def quantization_builder(n: common.BaseNode,
                          qat_config: QATConfig,
                          kernel_attr: str = None,
-                         ) -> Tuple[Dict[str, BasePytorchQATTrainableQuantizer],
-                                    List[BasePytorchQATTrainableQuantizer]]:
+                         ) -> Tuple[Dict[str, BasePytorchQATWeightTrainableQuantizer],
+                                    List[BasePytorchActivationTrainableQuantizer]]:
     """
     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,
-                                                        BasePytorchQATTrainableQuantizer)
+                                                        BasePytorchQATWeightTrainableQuantizer)
 
         weight_quantizers.update({kernel_attr: quantizer_class(get_trainable_quantizer_weights_config(n,
                                                                                                attr_name=kernel_attr,
@@ -95,7 +96,7 @@ def quantization_builder(n: common.BaseNode,
         quantizer_class = get_trainable_quantizer_class(QuantizationTarget.Activation,
                                                         qat_config.activation_training_method,
                                                         quant_method,
-                                                        BasePytorchQATTrainableQuantizer)
+                                                        BasePytorchActivationTrainableQuantizer)
 
         activation_quantizers = [quantizer_class(get_trainable_quantizer_activation_config(n, aq_cand),
                                                  **qat_config.activation_quantizer_params_override)]

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

@@ -18,28 +18,27 @@ import numpy as np
 import torch
 import torch.nn as nn
 
-from model_compression_toolkit.qat import TrainingMethod
 from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
 from mct_quantizers import PytorchQuantizationWrapper
 from model_compression_toolkit.qat.common import THRESHOLD_TENSOR
 from model_compression_toolkit import constants as C
-from model_compression_toolkit.qat.pytorch.quantizer.base_pytorch_qat_quantizer import BasePytorchQATTrainableQuantizer
+from model_compression_toolkit.qat.pytorch.quantizer.base_pytorch_qat_weight_quantizer import BasePytorchQATWeightTrainableQuantizer
 from mct_quantizers.common.base_inferable_quantizer import mark_quantizer, QuantizationTarget
 
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
-from model_compression_toolkit.qat.pytorch.quantizer.quantizer_utils import ste_round, ste_clip, symmetric_quantizer
+from model_compression_toolkit.trainable_infrastructure import TrainingMethod
+from model_compression_toolkit.trainable_infrastructure.pytorch.quantizer_utils import ste_round, ste_clip
 from mct_quantizers.pytorch.quantizers import \
-    WeightsPOTInferableQuantizer, WeightsSymmetricInferableQuantizer, ActivationPOTInferableQuantizer, \
-    ActivationSymmetricInferableQuantizer
+    WeightsPOTInferableQuantizer, WeightsSymmetricInferableQuantizer
 from model_compression_toolkit.trainable_infrastructure.common.trainable_quantizer_config import \
-    TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig
+    TrainableQuantizerWeightsConfig
 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(BasePytorchQATTrainableQuantizer):
+class STEWeightQATQuantizer(BasePytorchQATWeightTrainableQuantizer):
     """
     Trainable constrained quantizer to quantize a layer weights.
     """
@@ -127,80 +126,3 @@ class STEWeightQATQuantizer(BasePytorchQATTrainableQuantizer):
                                                       threshold=np_threshold.tolist(),
                                                       per_channel=self.quantization_config.weights_per_channel_threshold,
                                                       channel_axis=self.quantization_config.weights_channels_axis)
-
-
-
-@mark_quantizer(quantization_target=QuantizationTarget.Activation,
-                quantization_method=[QuantizationMethod.POWER_OF_TWO, QuantizationMethod.SYMMETRIC],
-                identifier=TrainingMethod.STE)
-class STEActivationQATQuantizer(BasePytorchQATTrainableQuantizer):
-    """
-    Trainable constrained quantizer to quantize a layer activations.
-    """
-
-    def __init__(self, quantization_config: TrainableQuantizerActivationConfig):
-        """
-        Initialize a STEActivationQATQuantizer object with parameters to use
-        for symmetric or power of two 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.sign = quantization_config.activation_quantization_params['is_signed']
-        np_threshold_values = quantization_config.activation_quantization_params[C.THRESHOLD]
-        self.threshold_tensor = torch.Tensor([np_threshold_values])
-        self.num_bits = quantization_config.activation_n_bits
-
-    def initialize_quantization(self,
-                                tensor_shape: torch.Size,
-                                name: str,
-                                layer: PytorchQuantizationWrapper):
-        """
-        Add quantizer parameters to the quantizer parameters dictionary
-
-        Args:
-            tensor_shape: tensor shape of the quantized tensor.
-            name: Tensor name.
-            layer: Layer to quantize.
-        """
-        layer.register_parameter(name, nn.Parameter(to_torch_tensor(self.threshold_tensor), requires_grad=True))
-
-        # save the quantizer added parameters for later calculations
-        self.add_quantizer_variable(THRESHOLD_TENSOR, layer.get_parameter(name), VariableGroup.QPARAMS)
-
-    def __call__(self,
-                 inputs: torch.Tensor,
-                 training: bool = True) -> torch.Tensor:
-        """
-        Quantize a tensor.
-        Args:
-            inputs: Input tensor to quantize.
-            training: Whether the graph is in training mode.
-
-        Returns:
-            The quantized tensor.
-        """
-
-        _t = self.get_quantizer_variable(THRESHOLD_TENSOR)
-        q_tensor = symmetric_quantizer(inputs, _t, self.num_bits, sign=self.sign)
-        return q_tensor
-
-    def convert2inferable(self) -> Union[ActivationPOTInferableQuantizer, ActivationSymmetricInferableQuantizer]:
-        """
-        Convert quantizer to inferable quantizer.
-
-        Returns:
-            A pytorch inferable quanizer object.
-        """
-        np_threshold = self.get_quantizer_variable(THRESHOLD_TENSOR).cpu().detach().numpy()
-        if self.power_of_two:
-            pot_threshold = np.power(2.0, np.ceil(np.log2(np_threshold)))
-            return ActivationPOTInferableQuantizer(num_bits=self.num_bits,
-                                                   threshold=pot_threshold.tolist(),
-                                                   signed=self.sign)
-        else:
-            return ActivationSymmetricInferableQuantizer(num_bits=self.num_bits,
-                                                         threshold=np_threshold.tolist(),
-                                                         signed=self.sign)

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

@@ -20,26 +20,25 @@ from torch import Tensor
 from model_compression_toolkit.constants import RANGE_MAX, RANGE_MIN
 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, PytorchQuantizationWrapper
-from model_compression_toolkit import constants as C
 
-from model_compression_toolkit.qat.pytorch.quantizer.base_pytorch_qat_quantizer import BasePytorchQATTrainableQuantizer
+from model_compression_toolkit.qat.pytorch.quantizer.base_pytorch_qat_weight_quantizer import BasePytorchQATWeightTrainableQuantizer
 from mct_quantizers import mark_quantizer
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
-from model_compression_toolkit.qat.pytorch.quantizer.quantizer_utils import uniform_quantizer
+from model_compression_toolkit.trainable_infrastructure import TrainingMethod
+from model_compression_toolkit.trainable_infrastructure.pytorch.quantizer_utils import uniform_quantizer
 from mct_quantizers.pytorch.quantizers import \
-    WeightsUniformInferableQuantizer, ActivationUniformInferableQuantizer
+    WeightsUniformInferableQuantizer
 from model_compression_toolkit.trainable_infrastructure.common.trainable_quantizer_config import \
-    TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig
+    TrainableQuantizerWeightsConfig
 from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
 @mark_quantizer(quantization_target=QuantizationTarget.Weights,
                 quantization_method=[QuantizationMethod.UNIFORM],
                 identifier=TrainingMethod.STE)
-class STEUniformWeightQATQuantizer(BasePytorchQATTrainableQuantizer):
+class STEUniformWeightQATQuantizer(BasePytorchQATWeightTrainableQuantizer):
     """
     Trainable constrained quantizer to quantize a layer inputs.
     """
@@ -66,7 +65,6 @@ class STEUniformWeightQATQuantizer(BasePytorchQATTrainableQuantizer):
                               [-1]) if self.quantization_config.weights_per_channel_threshold else float(
             self.min_values)
 
-
     def initialize_quantization(self,
                                 tensor_shape: torch.Size,
                                 name: str,
@@ -88,7 +86,6 @@ class STEUniformWeightQATQuantizer(BasePytorchQATTrainableQuantizer):
         self.add_quantizer_variable(FQ_MIN, layer.get_parameter(name+"_"+FQ_MIN), VariableGroup.QPARAMS)
         self.add_quantizer_variable(FQ_MAX, layer.get_parameter(name+"_"+FQ_MAX), VariableGroup.QPARAMS)
 
-
     def __call__(self,
                  inputs: nn.Parameter,
                  training: bool) -> Tensor:
@@ -117,79 +114,3 @@ class STEUniformWeightQATQuantizer(BasePytorchQATTrainableQuantizer):
                                                 max_range=_max.tolist(),
                                                 per_channel=self.quantization_config.weights_per_channel_threshold,
                                                 channel_axis=self.quantization_config.weights_channels_axis)
-
-
-@mark_quantizer(quantization_target=QuantizationTarget.Activation,
-                quantization_method=[QuantizationMethod.UNIFORM],
-                identifier=TrainingMethod.STE)
-class STEUniformActivationQATQuantizer(BasePytorchQATTrainableQuantizer):
-    """
-    Trainable constrained quantizer to quantize a layer activations.
-    """
-
-    def __init__(self, quantization_config: TrainableQuantizerActivationConfig):
-        """
-        Initialize a STEUniformActivationQATQuantizer object with parameters to use
-        for uniform quantization.
-
-        Args:
-            quantization_config: trainable quantizer config class
-        """
-        super().__init__(quantization_config)
-
-        np_min_range = quantization_config.activation_quantization_params[C.RANGE_MIN]
-        np_max_range = quantization_config.activation_quantization_params[C.RANGE_MAX]
-        self.min_range_tensor = torch.Tensor([np_min_range])
-        self.max_range_tensor = torch.Tensor([np_max_range])
-        self.num_bits = quantization_config.activation_n_bits
-
-    def initialize_quantization(self,
-                                tensor_shape: torch.Size,
-                                name: str,
-                                layer: PytorchQuantizationWrapper):
-        """
-        Add quantizer parameters to the quantizer parameters dictionary
-
-        Args:
-            tensor_shape: tensor shape of the quantized tensor.
-            name: Tensor name.
-            layer: Layer to quantize.
-        """
-        layer.register_parameter(name+"_"+FQ_MIN, nn.Parameter(to_torch_tensor(self.min_range_tensor), requires_grad=True))
-        layer.register_parameter(name+"_"+FQ_MAX, nn.Parameter(to_torch_tensor(self.max_range_tensor), requires_grad=True))
-
-        # Save the quantizer parameters for later calculations
-        self.add_quantizer_variable(FQ_MIN, layer.get_parameter(name+"_"+FQ_MIN), VariableGroup.QPARAMS)
-        self.add_quantizer_variable(FQ_MAX, layer.get_parameter(name+"_"+FQ_MAX), VariableGroup.QPARAMS)
-
-    def __call__(self,
-                 inputs: torch.Tensor,
-                 training: bool = True) -> torch.Tensor:
-        """
-        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 = uniform_quantizer(inputs, _min, _max, self.num_bits)
-        return q_tensor
-
-    def convert2inferable(self) -> ActivationUniformInferableQuantizer:
-        """
-        Convert quantizer to inferable quantizer.
-
-        Returns:
-            A pytorch inferable quanizer object.
-        """
-        _min = self.get_quantizer_variable(FQ_MIN).cpu().detach().numpy()
-        _max = self.get_quantizer_variable(FQ_MAX).cpu().detach().numpy()
-
-        return ActivationUniformInferableQuantizer(num_bits=self.num_bits,
-                                                   min_range=_min.tolist(),
-                                                   max_range=_max.tolist())

model_compression_toolkit/trainable_infrastructure/__init__.py

@@ -14,6 +14,12 @@
 # ==============================================================================
 
 from model_compression_toolkit.trainable_infrastructure.common.trainable_quantizer_config import TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig
-from model_compression_toolkit.trainable_infrastructure.keras.base_keras_quantizer import BaseKerasTrainableQuantizer
-from model_compression_toolkit.trainable_infrastructure.pytorch.base_pytorch_quantizer import BasePytorchTrainableQuantizer
-from model_compression_toolkit.trainable_infrastructure.keras.quantize_wrapper import KerasTrainableQuantizationWrapper
+from model_compression_toolkit.trainable_infrastructure.common.training_method import TrainingMethod
+from model_compression_toolkit.verify_packages import FOUND_TORCH, FOUND_TF
+if FOUND_TF:
+    from model_compression_toolkit.trainable_infrastructure.keras.base_keras_quantizer import BaseKerasTrainableQuantizer
+    from model_compression_toolkit.trainable_infrastructure.keras.quantize_wrapper import KerasTrainableQuantizationWrapper
+
+if FOUND_TORCH:
+    from model_compression_toolkit.trainable_infrastructure.pytorch.base_pytorch_quantizer import BasePytorchTrainableQuantizer
+    from model_compression_toolkit.trainable_infrastructure.pytorch.activation_quantizers import *

model_compression_toolkit/trainable_infrastructure/common/base_trainable_quantizer.py

@@ -12,12 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Union, List, Any
 from inspect import signature
 
-from model_compression_toolkit.core import common
 from model_compression_toolkit.logger import Logger
 
 from mct_quantizers.common.base_inferable_quantizer import BaseInferableQuantizer, \
@@ -31,6 +30,7 @@ from mct_quantizers.common.constants import QUANTIZATION_METHOD, \
 VAR = 'var'
 GROUP = 'group'
 
+
 class VariableGroup(Enum):
     """
     An enum for choosing trainable variable group
@@ -41,7 +41,7 @@ class VariableGroup(Enum):
     QPARAMS = 1
 
 
-class BaseTrainableQuantizer(BaseInferableQuantizer):
+class BaseTrainableQuantizer(BaseInferableQuantizer, ABC):
     def __init__(self,
                  quantization_config: Union[TrainableQuantizerActivationConfig, TrainableQuantizerWeightsConfig]):
         """
@@ -90,6 +90,7 @@ class BaseTrainableQuantizer(BaseInferableQuantizer):
     def get_sig(cls):
         return signature(cls)
 
+    @abstractmethod
     def initialize_quantization(self,
                                 tensor_shape,
                                 name: str,
@@ -105,8 +106,9 @@ class BaseTrainableQuantizer(BaseInferableQuantizer):
         Returns: None
 
         """
-        raise NotImplemented  # pragma: no cover
+        raise NotImplementedError()  # pragma: no cover
 
+    @abstractmethod
     def __call__(self,
                  input2quantize,
                  training: bool):
@@ -120,7 +122,7 @@ class BaseTrainableQuantizer(BaseInferableQuantizer):
         Returns:
             The quantized tensor.
         """
-        raise NotImplemented  # pragma: no cover
+        raise NotImplementedError()  # pragma: no cover
 
     def activation_quantization(self) -> bool:
         """
@@ -162,7 +164,7 @@ class BaseTrainableQuantizer(BaseInferableQuantizer):
         Returns:
             BaseInferableQuantizer object.
         """
-        raise NotImplemented  # pragma: no cover
+        raise NotImplementedError()  # pragma: no cover
 
     def add_quantizer_variable(self, name: str, variable: Any, group: VariableGroup = VariableGroup.WEIGHTS):
         """
@@ -185,7 +187,6 @@ class BaseTrainableQuantizer(BaseInferableQuantizer):
         else:
             Logger.critical(f"Variable '{name}' does not exist in quantizer parameters.") # pragma: no cover
 
-
     @abstractmethod
     def get_trainable_variables(self, group: VariableGroup) -> List[Any]:
         """
@@ -197,4 +198,4 @@ class BaseTrainableQuantizer(BaseInferableQuantizer):
         Returns:
             List of trainable variables
         """
-        raise NotImplemented  # pragma: no cover
+        raise NotImplementedError()  # pragma: no cover

model_compression_toolkit/trainable_infrastructure/common/training_method.py

@@ -0,0 +1,31 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from enum import Enum
+
+
+class TrainingMethod(Enum):
+    """
+    An enum for selecting a training method
+
+    STE - Standard straight-through estimator. Includes PowerOfTwo, symmetric & uniform quantizers
+
+    DQA -  DNN Quantization with Attention. Includes a smooth quantization introduces by DQA method
+
+    LSQ - Learned Step size Quantization. Includes PowerOfTwo, symmetric & uniform quantizers: https://arxiv.org/pdf/1902.08153.pdf
+
+    """
+    STE = "STE",
+    DQA = "DQA",
+    LSQ = "LSQ"

model_compression_toolkit/trainable_infrastructure/keras/base_keras_quantizer.py

@@ -81,8 +81,8 @@ if FOUND_TF:
 
 
 else:
-    class BaseKerasTrainableQuantizer(BaseTrainableQuantizer):
+    class BaseKerasTrainableQuantizer(BaseTrainableQuantizer):     # pragma: no cover
         def __init__(self, *args, **kwargs):
             Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
                             "BaseKerasTrainableQuantizer. The 'tensorflow' package is missing "
-                            "or is installed with a version higher than 2.15.")  # pragma: no cover
+                            "or is installed with a version higher than 2.15.")

model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py

@@ -101,7 +101,7 @@ if FOUND_TF:
             return inferable_quantizers_wrapper
 
 else:
-    class KerasTrainableQuantizationWrapper:
+    class KerasTrainableQuantizationWrapper:    # pragma: no cover
         def __init__(self, *args, **kwargs):
             """
             Keras Quantization Wrapper takes a keras layer and quantizers and infer a quantized layer.
@@ -112,4 +112,4 @@ else:
             """
             Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
                             "KerasTrainableQuantizationWrapper. The 'tensorflow' package is missing "
-                            "or is installed with a version higher than 2.15.")  # pragma: no cover
+                            "or is installed with a version higher than 2.15.")

model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from .base_activation_quantizer import BasePytorchActivationTrainableQuantizer
+from .ste.symmetric_ste import STESymmetricActivationTrainableQuantizer
+from .ste.uniform_ste import STEUniformActivationTrainableQuantizer
+from .lsq.symmetric_lsq import LSQSymmetricActivationTrainableQuantizer
+from .lsq.uniform_lsq import LSQUniformActivationTrainableQuantizer

model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/base_activation_quantizer.py

@@ -0,0 +1,22 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from abc import ABC
+
+from model_compression_toolkit.trainable_infrastructure import BasePytorchTrainableQuantizer
+
+
+class BasePytorchActivationTrainableQuantizer(BasePytorchTrainableQuantizer, ABC):
+    """ Base class for PyTorch trainable activation quantizers """
+    pass

model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/lsq/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/lsq/symmetric_lsq.py

@@ -0,0 +1,111 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from typing import Union
+
+import numpy as np
+import torch
+from torch import nn
+
+from mct_quantizers import mark_quantizer, QuantizationTarget, QuantizationMethod, PytorchQuantizationWrapper
+from mct_quantizers.pytorch.quantizers import ActivationPOTInferableQuantizer, ActivationSymmetricInferableQuantizer
+from model_compression_toolkit import constants as C
+from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
+from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerActivationConfig, TrainingMethod
+from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
+from model_compression_toolkit.trainable_infrastructure.common.constants import THRESHOLD_TENSOR
+from model_compression_toolkit.trainable_infrastructure.pytorch.activation_quantizers import \
+    BasePytorchActivationTrainableQuantizer
+from model_compression_toolkit.trainable_infrastructure.pytorch.quantizer_utils import symmetric_lsq_quantizer
+
+
+# moved (and renamed) from model_compression_toolkit/qat/pytorch/quantizer/lsq/symmetric_lsq.py
+@mark_quantizer(quantization_target=QuantizationTarget.Activation,
+                quantization_method=[QuantizationMethod.POWER_OF_TWO, QuantizationMethod.SYMMETRIC],
+                identifier=TrainingMethod.LSQ)
+class LSQSymmetricActivationTrainableQuantizer(BasePytorchActivationTrainableQuantizer):
+    """
+    Trainable constrained quantizer to quantize layer activations.
+    """
+
+    def __init__(self, quantization_config: TrainableQuantizerActivationConfig):
+        """
+        Initialize a LSQSymmetricActivationTrainableQuantizer object with parameters to use
+        for symmetric or power of two 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.sign = quantization_config.activation_quantization_params['is_signed']
+        self.threshold_values = np.array([quantization_config.activation_quantization_params[C.THRESHOLD]])
+        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
+
+    def initialize_quantization(self,
+                                tensor_shape: torch.Size,
+                                name: str,
+                                layer: PytorchQuantizationWrapper):
+        """
+        Add quantizer parameters to the quantizer parameters dictionary
+
+        Args:
+            tensor_shape: tensor shape of the quantized tensor.
+            name: Tensor name.
+            layer: Layer to quantize.
+        """
+        layer.register_parameter(name, nn.Parameter(to_torch_tensor(self.threshold_values), requires_grad=True))
+
+        # save the quantizer added parameters for later calculations
+        self.add_quantizer_variable(THRESHOLD_TENSOR, layer.get_parameter(name), VariableGroup.QPARAMS)
+
+    def __call__(self,
+                 inputs: torch.Tensor,
+                 training: bool = True) -> torch.Tensor:
+        """
+        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)
+        inputs_quantized = symmetric_lsq_quantizer(inputs, thresholds, self.num_bits, self.sign, self.min_int, self.max_int, scale_factor)
+        return inputs_quantized
+
+    def convert2inferable(self) -> Union[ActivationPOTInferableQuantizer, ActivationSymmetricInferableQuantizer]:
+        """
+        Convert quantizer to inferable quantizer.
+
+        Returns:
+            A pytorch inferable quanizer object.
+        """
+        threshold_values = self.get_quantizer_variable(THRESHOLD_TENSOR).cpu().detach().numpy()
+        if self.power_of_two:
+            pot_threshold = np.power(2.0, np.ceil(np.log2(threshold_values)))
+            return ActivationPOTInferableQuantizer(num_bits=self.num_bits,
+                                                   threshold=pot_threshold.tolist(),
+                                                   signed=self.sign)
+        else:
+            return ActivationSymmetricInferableQuantizer(num_bits=self.num_bits,
+                                                         threshold=threshold_values.tolist(),
+                                                         signed=self.sign)