mct-nightly 1.10.0.20231018.post358__py3-none-any.whl → 1.10.0.20231019.post424__py3-none-any.whl

{mct_nightly-1.10.0.20231018.post358.dist-info → mct_nightly-1.10.0.20231019.post424.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 1.10.0.20231018.post358
+Version: 1.10.0.20231019.post424
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct_nightly-1.10.0.20231018.post358.dist-info → mct_nightly-1.10.0.20231019.post424.dist-info}/RECORD

@@ -373,10 +373,13 @@ model_compression_toolkit/qat/keras/quantizer/ste_rounding/symmetric_ste.py,sha2
 model_compression_toolkit/qat/keras/quantizer/ste_rounding/uniform_ste.py,sha256=EED6LfqhX_OhDRJ9e4GwbpgNC9vq7hoXyJS2VPvG2qc,10789
 model_compression_toolkit/qat/pytorch/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/qat/pytorch/quantization_facade.py,sha256=LGIk_nbcQEOhSN8fe56u6QGGUtGOVhF7tlCwW2G-Ig4,12549
-model_compression_toolkit/qat/pytorch/quantizer/__init__.py,sha256=R4vwVcbg6QprCTNzibyF9PtbKKKBsfu9ffypKDNscJQ,859
+model_compression_toolkit/qat/pytorch/quantizer/__init__.py,sha256=xYa4C8pr9cG1f3mQQcBXO_u3IdJN-zl7leZxuXDs86w,1003
 model_compression_toolkit/qat/pytorch/quantizer/base_pytorch_qat_quantizer.py,sha256=FnhuFCuQoSf78FM1z1UZgXXd3k-mKSM7i9dYOuJUmeA,2213
 model_compression_toolkit/qat/pytorch/quantizer/quantization_builder.py,sha256=GOYRDXvQSGe_iUFVmvDy5BqC952hu_-rQO06n8QCyw0,5491
-model_compression_toolkit/qat/pytorch/quantizer/quantizer_utils.py,sha256=5XswoF-5aaFangkHssWKAQTsk6lf_zzndzfCsBWBVMs,5004
+model_compression_toolkit/qat/pytorch/quantizer/quantizer_utils.py,sha256=nO7IrDRo5b9Asf21WJacE4vf5voD3UzF_oGjBoGusD4,5335
+model_compression_toolkit/qat/pytorch/quantizer/lsq/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
+model_compression_toolkit/qat/pytorch/quantizer/lsq/symmetric_lsq.py,sha256=HihuaMi0P0OQkNZlZAE-QeYlK_4AqcDKV6N405SdgI0,10712
+model_compression_toolkit/qat/pytorch/quantizer/lsq/uniform_lsq.py,sha256=fqAI151SUo9OYN4UvQJIcdB6p1r7HVeteqPqHGxv-tI,10355
 model_compression_toolkit/qat/pytorch/quantizer/ste_rounding/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/qat/pytorch/quantizer/ste_rounding/symmetric_ste.py,sha256=4xmLmg7yN2A7iKnifwkWddgJTWMUiIjFilIuorJeK1A,9657
 model_compression_toolkit/qat/pytorch/quantizer/ste_rounding/uniform_ste.py,sha256=HshW016iVAMx7iMkUwlONN2P3K4XgDIu-2AnJnBVSGo,8778
@@ -445,8 +448,8 @@ model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py,sha
 model_compression_toolkit/trainable_infrastructure/keras/quantizer_utils.py,sha256=MVwXNymmFRB2NXIBx4e2mdJ1RfoHxRPYRgjb1MQP5kY,1797
 model_compression_toolkit/trainable_infrastructure/pytorch/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py,sha256=SbvRlIdE32PEBsINt1bhSqvrKL_zbM9V-aeSkOn-sw4,3083
-mct_nightly-1.10.0.20231018.post358.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-1.10.0.20231018.post358.dist-info/METADATA,sha256=mVpOJhL0dYj3nqrs9owiW0nPowR7YwjXGpEF7c0vt2E,16303
-mct_nightly-1.10.0.20231018.post358.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-mct_nightly-1.10.0.20231018.post358.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-1.10.0.20231018.post358.dist-info/RECORD,,
+mct_nightly-1.10.0.20231019.post424.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-1.10.0.20231019.post424.dist-info/METADATA,sha256=7KfKPP_oFOaieSu_3UTKRv1F3HaR5jiwMOP1DiBAiIg,16303
+mct_nightly-1.10.0.20231019.post424.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+mct_nightly-1.10.0.20231019.post424.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-1.10.0.20231019.post424.dist-info/RECORD,,

model_compression_toolkit/qat/pytorch/quantizer/__init__.py

@@ -14,4 +14,6 @@
 # ==============================================================================
 
 import model_compression_toolkit.qat.pytorch.quantizer.ste_rounding.symmetric_ste
-import model_compression_toolkit.qat.pytorch.quantizer.ste_rounding.uniform_ste
+import model_compression_toolkit.qat.pytorch.quantizer.ste_rounding.uniform_ste
+import model_compression_toolkit.qat.pytorch.quantizer.lsq.symmetric_lsq
+import model_compression_toolkit.qat.pytorch.quantizer.lsq.uniform_lsq

model_compression_toolkit/qat/pytorch/quantizer/lsq/__init__.py

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

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

@@ -0,0 +1,228 @@
+# Copyright 2023 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from typing import Union
+
+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 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, grad_scale
+from mct_quantizers.pytorch.quantizers import \
+    WeightsPOTInferableQuantizer, WeightsSymmetricInferableQuantizer, ActivationPOTInferableQuantizer, \
+    ActivationSymmetricInferableQuantizer
+from model_compression_toolkit.trainable_infrastructure.common.trainable_quantizer_config import \
+    TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig
+from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
+
+
+def symmetric_lsq_quantizer(x: nn.Parameter,
+                          thresholds: nn.Parameter,
+                          num_bits: int,
+                          sign: bool,
+                          min_int: int,
+                          max_int: int,
+                          scale_factor: float) -> Union[nn.Parameter, torch.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 = torch.clip(rounded, min=min_int, max=max_int)
+    quantized = delta_scaled * clipped
+    return quantized
+
+
+@mark_quantizer(quantization_target=QuantizationTarget.Weights,
+                quantization_method=[QuantizationMethod.POWER_OF_TWO, QuantizationMethod.SYMMETRIC],
+                identifier=TrainingMethod.LSQ)
+class LSQWeightQATQuantizer(BasePytorchQATTrainableQuantizer):
+    """
+    Trainable constrained quantizer to quantize layer's weights.
+    """
+
+    def __init__(self, quantization_config: TrainableQuantizerWeightsConfig):
+        """
+        Initialize a LSQWeightQATQuantizer 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.weights_quantization_method == QuantizationMethod.POWER_OF_TWO
+        self.threshold_values = np.array(quantization_config.weights_quantization_params[C.THRESHOLD])
+        if self.power_of_two:
+            self.threshold_values = np.power(2.0, np.ceil(np.log2(np.maximum(self.threshold_values, C.MIN_THRESHOLD))))
+        self.num_bits = self.quantization_config.weights_n_bits
+        n_pos_bits = self.num_bits - int(C.WEIGHTS_SIGNED)
+        self.min_int = -int(C.WEIGHTS_SIGNED) * (2 ** n_pos_bits)
+        self.max_int = 2 ** n_pos_bits - 1
+        self.scale_factor = 1.0 / np.sqrt(self.max_int * self.threshold_values.size)
+
+    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.
+        """
+
+        # Add threshold variables to layer.
+        layer.register_parameter(name + "_" + THRESHOLD_TENSOR, 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 + "_" + THRESHOLD_TENSOR), VariableGroup.QPARAMS)
+
+
+    def __call__(self,
+                 inputs: nn.Parameter,
+                 training: bool) -> nn.Parameter:
+        """
+        Quantize a tensor
+        Args:
+            inputs: Input tensor to quantize.
+            training: whether in training mode or not
+        Returns:
+            quantized tensor
+        """
+        thresholds = self.get_quantizer_variable(THRESHOLD_TENSOR)
+        weight_quantized = symmetric_lsq_quantizer(inputs, thresholds, self.num_bits, C.WEIGHTS_SIGNED, self.min_int, self.max_int, self.scale_factor)
+        return weight_quantized
+
+    def convert2inferable(self) -> Union[WeightsPOTInferableQuantizer, WeightsSymmetricInferableQuantizer]:
+        """
+        Convert quantizer to inferable quantizer.
+
+        Returns:
+            A pytorch inferable quanizer object.
+        """
+        threshold_values = self.get_quantizer_variable(THRESHOLD_TENSOR).cpu().detach().numpy().flatten()
+        if self.power_of_two:
+            pot_threshold = 2 ** np.ceil(np.log2(threshold_values))
+            return WeightsPOTInferableQuantizer(num_bits=self.num_bits,
+                                                threshold=pot_threshold.tolist(),
+                                                per_channel=self.quantization_config.weights_per_channel_threshold,
+                                                channel_axis=self.quantization_config.weights_channels_axis)
+        else:
+            return WeightsSymmetricInferableQuantizer(num_bits=self.num_bits,
+                                                      threshold=threshold_values.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.LSQ)
+class LSQActivationQATQuantizer(BasePytorchQATTrainableQuantizer):
+    """
+    Trainable constrained quantizer to quantize layer activations.
+    """
+
+    def __init__(self, quantization_config: TrainableQuantizerActivationConfig):
+        """
+        Initialize a LSQActivationQATQuantizer 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)

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

@@ -0,0 +1,223 @@
+# Copyright 2023 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from typing import Union
+import numpy as np
+import torch
+import torch.nn as nn
+
+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 mct_quantizers import mark_quantizer
+from model_compression_toolkit.qat.pytorch.quantizer.quantizer_utils import ste_round, grad_scale
+from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
+from mct_quantizers.pytorch.quantizers import \
+    WeightsUniformInferableQuantizer, ActivationUniformInferableQuantizer
+from model_compression_toolkit.trainable_infrastructure.common.trainable_quantizer_config import \
+    TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig
+from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
+from model_compression_toolkit.qat.pytorch.quantizer.quantizer_utils import adjust_range_to_include_zero
+from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import fix_range_to_include_zero
+
+
+
+def uniform_lsq_quantizer(x: nn.Parameter,
+                      min_range: nn.Parameter,
+                      max_range: nn.Parameter,
+                      num_bits: int,
+                      min_int: int,
+                      max_int: int,
+                      scale_factor: float) -> Union[nn.Parameter, torch.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
+    """
+    a, b = adjust_range_to_include_zero(min_range, max_range, num_bits)
+    delta = (b - a) / (2 ** num_bits - 1)
+    delta_scaled = grad_scale(delta, scale_factor)
+    rounded = ste_round((x - a) / delta_scaled)
+    clipped = torch.clip(rounded, min=min_int, max=max_int)
+    quantized = delta_scaled * clipped + a
+    return quantized
+
+
+@mark_quantizer(quantization_target=QuantizationTarget.Weights,
+                quantization_method=[QuantizationMethod.UNIFORM],
+                identifier=TrainingMethod.LSQ)
+class LSQUniformWeightQATQuantizer(BasePytorchQATTrainableQuantizer):
+    """
+    Trainable constrained quantizer to quantize layer's weights.
+    """
+
+    def __init__(self, quantization_config: TrainableQuantizerWeightsConfig):
+        """
+        Initialize a LSQUniformWeightQATQuantizer object with parameters to use
+        for the quantization.
+
+        Args:
+            quantization_config: trainable quantizer config class
+        """
+        super().__init__(quantization_config)
+        self.num_bits = self.quantization_config.weights_n_bits
+        self.min_int = 0
+        self.max_int = 2 ** self.num_bits - 1
+        self.min_values = np.array(quantization_config.weights_quantization_params[RANGE_MIN])
+        self.max_values = np.array(quantization_config.weights_quantization_params[RANGE_MAX])
+        self.scale_factor = 1.0 / np.sqrt(self.max_int * self.min_values.size)
+
+    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.
+        """
+
+        # Add min and max variables to layer.
+        layer.register_parameter(name+"_"+FQ_MIN, nn.Parameter(to_torch_tensor(self.min_values), requires_grad=True))
+        layer.register_parameter(name+"_"+FQ_MAX, nn.Parameter(to_torch_tensor(self.max_values), 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: nn.Parameter,
+                 training: bool) -> torch.Tensor:
+        """
+        Quantize a tensor
+        Args:
+            inputs: Input tensor to quantize.
+            training: whether in training mode or not
+        Returns:
+            quantized tensor
+        """
+        min_range = self.get_quantizer_variable(FQ_MIN)
+        max_range = self.get_quantizer_variable(FQ_MAX)
+        weight_quantized = uniform_lsq_quantizer(inputs, min_range, max_range, self.num_bits, self.min_int, self.max_int, self.scale_factor)
+        return weight_quantized
+
+    def convert2inferable(self) -> WeightsUniformInferableQuantizer:
+        """
+        Convert quantizer to inferable quantizer.
+
+        Returns:
+            A pytorch inferable quanizer object.
+        """
+        min_range = self.get_quantizer_variable(FQ_MIN).cpu().detach().numpy()
+        max_range = self.get_quantizer_variable(FQ_MAX).cpu().detach().numpy()
+        min_range, max_range = fix_range_to_include_zero(min_range, max_range, self.num_bits)
+        return WeightsUniformInferableQuantizer(num_bits=self.num_bits,
+                                                min_range=min_range.tolist(),
+                                                max_range=max_range.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.LSQ)
+class LSQUniformActivationQATQuantizer(BasePytorchQATTrainableQuantizer):
+    """
+    Trainable constrained quantizer to quantize layer activations.
+    """
+
+    def __init__(self, quantization_config: TrainableQuantizerActivationConfig):
+        """
+        Initialize a LSQUniformActivationQATQuantizer object with parameters to use
+        for uniform quantization.
+
+        Args:
+            quantization_config: trainable quantizer config class
+        """
+        super().__init__(quantization_config)
+        self.num_bits = self.quantization_config.activation_n_bits
+        self.min_int = 0
+        self.max_int = 2 ** self.num_bits - 1
+        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]])
+
+    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), requires_grad=True))
+        layer.register_parameter(name+"_"+FQ_MAX, nn.Parameter(to_torch_tensor(self.max_range), 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_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)
+        inputs_quantized = uniform_lsq_quantizer(inputs, min_range, max_range, self.num_bits, self.min_int, self.max_int, scale_factor)
+        return inputs_quantized
+
+    def convert2inferable(self) -> ActivationUniformInferableQuantizer:
+        """
+        Convert quantizer to inferable quantizer.
+
+        Returns:
+            A pytorch inferable quanizer object.
+        """
+        min_range = self.get_quantizer_variable(FQ_MIN).cpu().detach().numpy()
+        max_range = self.get_quantizer_variable(FQ_MAX).cpu().detach().numpy()
+        min_range, max_range = fix_range_to_include_zero(min_range, max_range, self.num_bits)
+        return ActivationUniformInferableQuantizer(num_bits=self.num_bits,
+                                                   min_range=min_range.tolist(),
+                                                   max_range=max_range.tolist())

model_compression_toolkit/qat/pytorch/quantizer/quantizer_utils.py

@@ -40,9 +40,22 @@ def ste_clip(x: torch.Tensor, min_val=-1.0, max_val=1.0) -> torch.Tensor:
     return (torch.clip(x, min=min_val, max=max_val) - x).detach() + x
 
 
-def fix_range_to_include_zero(range_min: torch.Tensor,
-                              range_max: torch.Tensor,
-                              n_bits: int) -> Tuple[torch.Tensor, torch.Tensor]:
+def grad_scale(x: torch.Tensor, scale=1.0) -> torch.Tensor:
+    """
+    Gradient scale
+    Args:
+        x: input variable
+        scale: scale factor
+    Returns:
+        x in forward and x*scale in backward (for scaling the gradients).
+    """
+    x_scaled = x * scale
+    return (x - x_scaled).detach() + x_scaled
+
+
+def adjust_range_to_include_zero(range_min: torch.Tensor,
+                                 range_max: torch.Tensor,
+                                 n_bits: int) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Adjusting the quantization range to include representation of 0.0 in the quantization grid.
     If quantization per-channel, then range_min and range_max should be tensors in the specific shape that allows
@@ -120,7 +133,7 @@ def uniform_quantizer(tensor_data: torch.Tensor,
         Quantized data.
     """
     # adjusts the quantization range so the quantization grid includes zero.
-    a, b = fix_range_to_include_zero(range_min, range_max, n_bits)
+    a, b = adjust_range_to_include_zero(range_min, range_max, n_bits)
 
     # Compute the step size of quantized values.
     delta_tensor = (b - a) / (2 ** n_bits - 1)