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

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

@@ -0,0 +1,106 @@
+# 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.
+# ==============================================================================
+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 ActivationUniformInferableQuantizer
+from model_compression_toolkit import constants as C
+from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import fix_range_to_include_zero
+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 FQ_MIN, FQ_MAX
+from model_compression_toolkit.trainable_infrastructure.pytorch.activation_quantizers import \
+    BasePytorchActivationTrainableQuantizer
+from model_compression_toolkit.trainable_infrastructure.pytorch.quantizer_utils import uniform_lsq_quantizer
+
+
+# moved (and renamed) from model_compression_toolkit/qat/pytorch/quantizer/lsq/uniform_lsq.py
+@mark_quantizer(quantization_target=QuantizationTarget.Activation,
+                quantization_method=[QuantizationMethod.UNIFORM],
+                identifier=TrainingMethod.LSQ)
+class LSQUniformActivationTrainableQuantizer(BasePytorchActivationTrainableQuantizer):
+    """
+    Trainable constrained quantizer to quantize layer activations.
+    """
+
+    def __init__(self, quantization_config: TrainableQuantizerActivationConfig):
+        """
+        Initialize a LSQUniformActivationTrainableQuantizer 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/trainable_infrastructure/pytorch/activation_quantizers/ste/__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/ste/symmetric_ste.py

@@ -0,0 +1,108 @@
+# 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 TrainingMethod, TrainableQuantizerActivationConfig
+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_quantizer
+
+
+# moved (and renamed) from model_compression_toolkit/qat/pytorch/quantizer/ste_rounding/symmetric_ste.py
+@mark_quantizer(quantization_target=QuantizationTarget.Activation,
+                quantization_method=[QuantizationMethod.POWER_OF_TWO, QuantizationMethod.SYMMETRIC],
+                identifier=TrainingMethod.STE)
+class STESymmetricActivationTrainableQuantizer(BasePytorchActivationTrainableQuantizer):
+    """
+    Trainable constrained quantizer to quantize a layer activations.
+    """
+
+    def __init__(self, quantization_config: TrainableQuantizerActivationConfig):
+        """
+        Initialize a STESymmetricActivationTrainableQuantizer 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/trainable_infrastructure/pytorch/activation_quantizers/ste/uniform_ste.py

@@ -0,0 +1,105 @@
+# 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.
+# ==============================================================================
+import torch
+from torch import nn
+
+from mct_quantizers import mark_quantizer, QuantizationTarget, QuantizationMethod, PytorchQuantizationWrapper
+from mct_quantizers.pytorch.quantizers import ActivationUniformInferableQuantizer
+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 FQ_MIN, FQ_MAX
+from model_compression_toolkit.trainable_infrastructure.pytorch.activation_quantizers import \
+    BasePytorchActivationTrainableQuantizer
+from model_compression_toolkit.trainable_infrastructure.pytorch.quantizer_utils import uniform_quantizer
+
+
+# moved (and renamed) from model_compression_toolkit/qat/pytorch/quantizer/ste_rounding/uniform_ste.py
+@mark_quantizer(quantization_target=QuantizationTarget.Activation,
+                quantization_method=[QuantizationMethod.UNIFORM],
+                identifier=TrainingMethod.STE)
+class STEUniformActivationTrainableQuantizer(BasePytorchActivationTrainableQuantizer):
+    """
+    Trainable constrained quantizer to quantize a layer activations.
+    """
+
+    def __init__(self, quantization_config: TrainableQuantizerActivationConfig):
+        """
+        Initialize a STEUniformActivationTrainableQuantizer 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/pytorch/base_pytorch_quantizer.py

@@ -12,31 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Union, List
+from typing import List
+
+from abc import ABC
 
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import BaseTrainableQuantizer, VAR, GROUP
-from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig, \
-    TrainableQuantizerActivationConfig
 
 
 if FOUND_TORCH:
 
     import torch
 
-    class BasePytorchTrainableQuantizer(BaseTrainableQuantizer):
-        def __init__(self,
-                     quantization_config: Union[TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig]):
-            """
-            This class is a base Pytorch quantizer which validates the provided quantization config and defines an
-            abstract function which any quantizer needs to implement.
-
-            Args:
-                quantization_config: quantizer config class contains all the information about the quantizer configuration.
-            """
-            super().__init__(quantization_config)
+    class BasePytorchTrainableQuantizer(BaseTrainableQuantizer, ABC):
+        """
+        Base class for PyTorch trainable quantizers
+        """
 
         def get_trainable_variables(self, group: VariableGroup) -> List[torch.Tensor]:
             """

model_compression_toolkit/{qat/pytorch/quantizer → trainable_infrastructure/pytorch}/quantizer_utils.py

@@ -12,15 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Tuple
+from typing import Tuple, Union
 import torch
+from torch import nn
 
 
 def ste_round(x: torch.Tensor) -> torch.Tensor:
     """
     Calculate the rounded values of a tensor
+
     Args:
         x: input variable
+
     Returns:
         rounded value
     """
@@ -30,10 +33,12 @@ def ste_round(x: torch.Tensor) -> torch.Tensor:
 def ste_clip(x: torch.Tensor, min_val=-1.0, max_val=1.0) -> torch.Tensor:
     """
     Clip a variable between fixed values such that min_val<=output<=max_val
+
     Args:
         x: input variable
         min_val: minimum value for clipping
         max_val: maximum value for clipping
+
     Returns:
         clipped variable
     """
@@ -43,9 +48,11 @@ def ste_clip(x: torch.Tensor, min_val=-1.0, max_val=1.0) -> 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).
     """
@@ -60,11 +67,14 @@ def adjust_range_to_include_zero(range_min: 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
     quantization along the channel_axis.
+
     Args:
         range_min: min bound of the quantization range (before adjustment).
         range_max: max bound of the quantization range (before adjustment).
         n_bits: Number of bits to quantize the tensor.
-    Returns: adjusted quantization range
+
+    Returns:
+        adjusted quantization range
     """
     min_positive = range_min > 0
     max_negative = range_max < 0
@@ -89,11 +99,13 @@ def symmetric_quantizer(tensor_data: torch.Tensor,
     """
     Quantize a tensor according to the number of bits and threshold.
     Symmetric quantization.
+
     Args:
         tensor_data: Tensor values to quantize.
         threshold: threshold for quantization.
         n_bits: Number of bits to quantize the tensor.
         sign: sign of tensor_data
+
     Returns:
         Quantized data.
     """
@@ -124,11 +136,13 @@ def uniform_quantizer(tensor_data: torch.Tensor,
     """
     Quantize a tensor according to given range (min, max) and number of bits.
     Uniform quantization.
+
     Args:
         tensor_data: Tensor values to quantize.
         range_min: minimum bound of the range for quantization (or array of min values per channel).
         range_max: maximum bound of the range for quantization (or array of max values per channel).
         n_bits: Number of bits to quantize the tensor.
+
     Returns:
         Quantized data.
     """
@@ -147,3 +161,66 @@ def uniform_quantizer(tensor_data: torch.Tensor,
     # Quantize the data between min/max of quantization range.
     q = delta_tensor * clipped_tensor + a
     return q
+
+
+# moved from model_compression_toolkit/qat/pytorch/quantizer/lsq/symmetric_lsq.py
+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
+
+
+# moved from model_compression_toolkit/qat/pytorch/quantizer/lsq/uniform_lsq.py
+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