mct-nightly 1.7.1.31122022.post351__py3-none-any.whl → 1.8.0.1042023.post423__py3-none-any.whl

model_compression_toolkit/quantizers_infrastructure/trainable_infrastructure/common/get_quantizer_config.py

@@ -0,0 +1,116 @@
+# 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 List
+from model_compression_toolkit.core.common import BaseNode, Logger
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.trainable_quantizer_config import \
+    TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig, TrainableQuantizerCandidateConfig
+
+
+def get_trainable_quantizer_weights_config(
+        n: BaseNode,
+        weights_quantization_candidates: List[TrainableQuantizerCandidateConfig] = None
+) -> TrainableQuantizerWeightsConfig:
+    """
+    Returns the relevant configuration for weights trainable quantizer
+
+    Args:
+        n: BaseNode - the node to build a trainable quantizer from.
+        weights_quantization_candidates: A list of weights quantizer config candidates.
+
+    Returns:
+         TrainableQuantizerWeightsConfig: an object that contains the quantizer configuration
+    """
+    if n.final_weights_quantization_cfg is None:
+        Logger.error(f'Node must have final_weights_quantization_cfg in order to build quantizer configuration')  # pragma: no cover
+
+    final_cfg = n.final_weights_quantization_cfg
+    return TrainableQuantizerWeightsConfig(final_cfg.weights_quantization_method,
+                                           final_cfg.weights_n_bits,
+                                           final_cfg.weights_quantization_params,
+                                           final_cfg.enable_weights_quantization,
+                                           final_cfg.weights_channels_axis,
+                                           final_cfg.weights_per_channel_threshold,
+                                           final_cfg.min_threshold,
+                                           weights_quantization_candidates)
+
+
+def get_trainable_quantizer_activation_config(
+        n: BaseNode,
+        activation_quantization_candidates: List[TrainableQuantizerCandidateConfig] = None
+) -> TrainableQuantizerActivationConfig:
+    """
+    Returns configuration for activation trainable quantizer
+
+    Args:
+        n: BaseNode - the node to build a trainable quantizer from.
+        activation_quantization_candidates: A list of activation quantizer candidates config.
+
+    Returns:
+         TrainableQuantizerActivationConfig - an object that contains the quantizer configuration
+    """
+    if n.final_activation_quantization_cfg is None:
+        Logger.error(f'Node must have final_activation_quantization_cfg in order to build quantizer configuration')  # pragma: no cover
+
+    final_cfg = n.final_activation_quantization_cfg
+    return TrainableQuantizerActivationConfig(final_cfg.activation_quantization_method,
+                                              final_cfg.activation_n_bits,
+                                              final_cfg.activation_quantization_params,
+                                              final_cfg.enable_activation_quantization,
+                                              final_cfg.min_threshold,
+                                              activation_quantization_candidates)
+
+
+def get_trainable_quantizer_quantization_candidates(n: BaseNode):
+    """
+    Returns quantization configuration candidates for activation and weights trainable quantizer.
+    Checks that the candidates are compatible with trainable quantizer
+
+    Args:
+        n: BaseNode - the node to build a trainable quantizer from
+
+    Returns:
+         weights_quantization_candidates - A list of configuration candidates for weights
+         activation_quantization_candidates - A list of configuration candidates for activation
+    """
+    # all candidates must have the same weights quantization method
+    weights_quantization_methods = set([cfg.weights_quantization_cfg.weights_quantization_method for cfg in n.candidates_quantization_cfg])
+    if len(weights_quantization_methods) > 1:
+        Logger.error(f'Unsupported candidates_quantization_cfg with different weights quantization methods: {weights_quantization_methods}') # pragma: no cover
+
+    # all candidates must have the same activation quantization method
+    activation_quantization_methods = set([cfg.activation_quantization_cfg.activation_quantization_method for cfg in n.candidates_quantization_cfg])
+    if len(activation_quantization_methods) > 1:
+        Logger.error(f'Unsupported candidates_quantization_cfg with different activation quantization methods: {activation_quantization_methods}') # pragma: no cover
+
+    # get unique lists of candidates
+    unique_weights_candidates = n.get_unique_weights_candidates()
+    unique_activation_candidates = n.get_unique_activation_candidates()
+
+    # verify all the combinations of weights_n_bits and activation_n_bits are allowed
+    if len(n.candidates_quantization_cfg) != len(unique_weights_candidates) * len(unique_activation_candidates):
+        Logger.error(f'Unsupported candidates_quantization_cfg for a trainable quantizer,'
+                            f'it must contain all the combinations of (weights_n_bits X activations_n_bits)') # pragma: no cover
+
+    # generate list of weights quantizer candidates
+    weights_cfg_candidates = [TrainableQuantizerCandidateConfig(
+        cfg.weights_quantization_cfg.weights_n_bits,
+        cfg.weights_quantization_cfg.weights_quantization_params) for cfg in unique_weights_candidates]
+
+    # generate list of activation quantizer candidates
+    activation_cfg_candidates = [TrainableQuantizerCandidateConfig(
+        cfg.activation_quantization_cfg.activation_n_bits,
+        cfg.activation_quantization_cfg.activation_quantization_params) for cfg in unique_activation_candidates]
+
+    return weights_cfg_candidates, activation_cfg_candidates

model_compression_toolkit/quantizers_infrastructure/trainable_infrastructure/common/get_quantizers.py

@@ -0,0 +1,65 @@
+# 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
+
+from model_compression_toolkit.gptq import RoundingType
+from model_compression_toolkit import TrainingMethod
+from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.core.common.target_platform import QuantizationMethod
+from model_compression_toolkit.quantizers_infrastructure import QuantizationTarget
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.constants \
+    import QUANTIZATION_TARGET, QUANTIZATION_METHOD, QUANTIZER_TYPE
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.get_all_subclasses \
+    import get_all_subclasses
+
+
+def get_trainable_quantizer_class(quant_target: QuantizationTarget,
+                                  quantizer_type: Union[TrainingMethod, RoundingType],
+                                  quant_method: QuantizationMethod,
+                                  quantizer_base_class: type) -> type:
+    """
+    Searches for a trainable quantizer class that matches the requested QuantizationTarget and QuantizationMethod and
+    a task dedicated quantizer type. Exactly one class should be found.
+
+    Args:
+        quant_target: QuantizationTarget value which indicates what is the target for quantization to
+            use the quantizer for.
+        quantizer_type: The type of the quantizer (quantization technique).
+            This can differ, depending on the purpose the quantizer is for.
+        quant_method: A list of QuantizationMethod values to indicate all type of quantization methods that the
+            quantizer supports.
+        quantizer_base_class: A type of quantizer that the requested quantizer should inherit from.
+
+    Returns: A class of a quantizer that inherits from BaseKerasQATTrainableQuantizer.
+
+    """
+    qat_quantizer_classes = get_all_subclasses(quantizer_base_class)
+    if len(qat_quantizer_classes) == 0:
+        Logger.error(f"No quantizers were found that inherit from {quantizer_base_class}.")  # pragma: no cover
+
+    filtered_quantizers = list(filter(lambda q_class: getattr(q_class, QUANTIZATION_TARGET, None) is not None and
+                                                      getattr(q_class, QUANTIZATION_TARGET) == quant_target and
+                                                      getattr(q_class, QUANTIZATION_METHOD, None) is not None and
+                                                       quant_method in getattr(q_class, QUANTIZATION_METHOD, []) and
+                                                      getattr(q_class, QUANTIZER_TYPE, None) == quantizer_type,
+                                      qat_quantizer_classes))
+
+    if len(filtered_quantizers) != 1:
+        Logger.error(f"Found {len(filtered_quantizers)} quantizer for target {quant_target.value} "  # pragma: no cover
+                     f"that matches the requested quantization method {quant_method.name} and "
+                     f"quantizer type {quantizer_type.value} but there should be exactly one."
+                     f"The possible quantizers that were found are {filtered_quantizers}.")
+
+    return filtered_quantizers[0]

model_compression_toolkit/quantizers_infrastructure/trainable_infrastructure/common/quant_utils.py

@@ -0,0 +1,36 @@
+# 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 Tuple, List
+
+
+def get_threshold_reshape_shape(tensor_shape: Tuple, quant_axis: int, quant_axis_dim: int) -> List[int]:
+    """
+    Gets a shape that contains 1 in all axis except the quantization axis, to adjust the threshold tensor for
+    per-channel quantization.
+
+    Args:
+        tensor_shape: The shape of th
+
+        e tensor to be quantized.
+        quant_axis: The axis along which the quantization happens (usually the tensor's channel axis).
+        quant_axis_dim: The dimension of the quantization axis.
+
+    Returns: A shape to reshape the threshold tensor according to.
+
+    """
+    n_axis = len(tensor_shape)
+    quantization_axis = n_axis + quant_axis if quant_axis < 0 else quant_axis
+
+    return [quant_axis_dim if i == quantization_axis else 1 for i in range(n_axis)]

model_compression_toolkit/quantizers_infrastructure/trainable_infrastructure/common/trainable_quantizer_config.py

@@ -0,0 +1,97 @@
+# 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 abc import ABC
+from typing import Dict, List
+from model_compression_toolkit.core.common.target_platform import QuantizationMethod
+
+
+class TrainableQuantizerCandidateConfig:
+
+    def __init__(self,
+                 n_bits: int,
+                 quantization_params: Dict,
+                 ):
+        """
+        Class for representing candidates of quantization configurations for trainable quantizer.
+        It can be used for weights and activation quantization configuration.
+
+        Args:
+            n_bits (int): Number of bits to use for quantization.
+            quantization_params (Dict): Dictionary that contains quantization params.
+        """
+
+        self.n_bits = n_bits
+        self.quantization_params = quantization_params
+
+
+class TrainableQuantizerActivationConfig:
+
+    def __init__(self,
+                 activation_quantization_method: QuantizationMethod,
+                 activation_n_bits: int,
+                 activation_quantization_params: Dict,
+                 enable_activation_quantization: bool,
+                 min_threshold: float,
+                 activation_quantization_candidates: List[TrainableQuantizerCandidateConfig] = None,
+                 ):
+        """
+        Attributes for configuring activations trainable quantizer.
+
+        Args:
+            activation_quantization_method (QuantizationMethod): Which method to use from QuantizationMethod for activation quantization.
+            activation_n_bits (int): Number of bits to quantize the activations.
+            activation_quantization_params (Dict): Dictionary that contains activation quantization params.
+            enable_activation_quantization (bool): Whether to quantize the layer's activations or not.
+            min_threshold (float): Minimum threshold to use during thresholds selection.
+        """
+        self.activation_quantization_method = activation_quantization_method
+        self.activation_n_bits = activation_n_bits
+        self.activation_quantization_params = activation_quantization_params
+        self.enable_activation_quantization = enable_activation_quantization
+        self.min_threshold = min_threshold
+        self.activation_bits_candidates = activation_quantization_candidates
+
+
+class TrainableQuantizerWeightsConfig:
+    def __init__(self,
+                 weights_quantization_method: QuantizationMethod,
+                 weights_n_bits: int,
+                 weights_quantization_params: Dict,
+                 enable_weights_quantization: bool,
+                 weights_channels_axis: int,
+                 weights_per_channel_threshold: bool,
+                 min_threshold: float,
+                 weights_quantization_candidates: List[TrainableQuantizerCandidateConfig] = None,
+                 ):
+        """
+        Attributes for configuring weights trainable quantizer.
+
+        Args:
+            weights_quantization_method (QuantizationMethod): Which method to use from QuantizationMethod for weights quantization.
+            weights_n_bits (int): Number of bits to quantize the coefficients.
+            weights_quantization_params (Dict): Dictionary that contains weights quantization params.
+            enable_weights_quantization (bool): Whether to quantize the layer's weights or not.
+            weights_channels_axis (int): Axis to quantize a node's kernel when quantizing per-channel.
+            weights_per_channel_threshold (bool): Whether to quantize the weights per-channel or not (per-tensor).
+            min_threshold (float): Minimum threshold to use during thresholds selection.
+        """
+        self.weights_quantization_method = weights_quantization_method
+        self.weights_n_bits = weights_n_bits
+        self.weights_quantization_params = weights_quantization_params
+        self.enable_weights_quantization = enable_weights_quantization
+        self.weights_channels_axis = weights_channels_axis
+        self.weights_per_channel_threshold = weights_per_channel_threshold
+        self.min_threshold = min_threshold
+        self.weights_bits_candidates = weights_quantization_candidates

model_compression_toolkit/quantizers_infrastructure/trainable_infrastructure/keras/__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/quantizers_infrastructure/trainable_infrastructure/keras/base_keras_quantizer.py

@@ -0,0 +1,90 @@
+# Copyright 2022 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 Dict, Any, Union, List
+
+from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.core.common.constants import FOUND_TF
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import BaseTrainableQuantizer, VAR, GROUP
+from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig, \
+    TrainableQuantizerActivationConfig
+
+if FOUND_TF:
+    QUANTIZATION_CONFIG = 'quantization_config'
+    from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.keras.config_serialization import config_serialization, \
+        config_deserialization
+    import tensorflow as tf
+
+    class BaseKerasTrainableQuantizer(BaseTrainableQuantizer):
+        def __init__(self,
+                     quantization_config: Union[TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig]):
+            """
+            This class is a base quantizer which validates provided quantization config and defines an abstract function which any quantizer needs to implement.
+            This class adds to the base quantizer a get_config and from_config functions to enable loading and saving the keras model.
+            Args:
+                quantization_config: quantizer config class contains all the information about a quantizer configuration.
+            """
+            super().__init__(quantization_config)
+
+        def get_config(self) -> Dict[str, Any]:
+            """
+
+            Returns: Configuration of BaseKerasQuantizer.
+
+            """
+            return {QUANTIZATION_CONFIG: config_serialization(self.quantization_config)}
+
+        @classmethod
+        def from_config(cls, config: dict):
+            """
+
+            Args:
+                config(dict): dictonory  of  BaseKerasQuantizer Configuration
+
+            Returns: A BaseKerasQuantizer
+
+            """
+            config = config.copy()
+            quantization_config = config_deserialization(config[QUANTIZATION_CONFIG])
+            # Note that a quantizer only receive quantization config and the rest of define hardcoded inside the speficie quantizer.
+            return cls(quantization_config=quantization_config)
+
+        def get_trainable_variables(self, group: VariableGroup) -> List[tf.Tensor]:
+            """
+            Get trainable parameters with specific group from quantizer
+
+            Args:
+                group: Enum of variable group
+
+            Returns:
+                List of trainable variables
+            """
+            quantizer_trainable = []
+            for name, parameter_dict in self.quantizer_parameters.items():
+                quantizer_parameter, parameter_group = parameter_dict[VAR], parameter_dict[GROUP]
+                if quantizer_parameter.trainable and parameter_group == group:
+                    quantizer_trainable.append(quantizer_parameter)
+            return quantizer_trainable
+
+
+else:
+    class BaseKerasTrainableQuantizer(BaseTrainableQuantizer):
+        def __init__(self,
+                     quantization_config: Union[TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig]):
+
+            super().__init__(quantization_config)
+            Logger.critical('Installing tensorflow and tensorflow_model_optimization is mandatory '
+                            'when using BaseKerasQuantizer. '
+                            'Could not find Tensorflow package.')  # pragma: no cover

model_compression_toolkit/quantizers_infrastructure/trainable_infrastructure/keras/config_serialization.py

@@ -0,0 +1,80 @@
+# Copyright 2022 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 copy
+
+from typing import Any, Union
+from enum import Enum
+
+from model_compression_toolkit.core.common.target_platform import QuantizationMethod
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.trainable_quantizer_config import \
+    TrainableQuantizerActivationConfig, TrainableQuantizerWeightsConfig
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common import constants as C
+
+
+def transform_enum(v: Any):
+    """
+    If an enum is received it value is return otherwise the input is returned.
+    Args:
+        v: Any type
+
+    Returns: Any
+
+    """
+    if isinstance(v, Enum):
+        return v.value
+    return v
+
+
+def config_serialization(quantization_config: Union[TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig]):
+    """
+    This function change trainable quantizer config to a dictionary
+    Args:
+        quantization_config: A TrainableQuantizerWeightsConfig or TrainableQuantizerActivationConfig for serialization
+
+    Returns: A config dictionary of quantizer config
+
+    """
+    config_data = {k: transform_enum(v) for k, v in quantization_config.__dict__.items()}
+    config_data[C.IS_WEIGHTS] = isinstance(quantization_config, TrainableQuantizerWeightsConfig)
+    config_data[C.IS_ACTIVATIONS] = isinstance(quantization_config, TrainableQuantizerActivationConfig)
+    return config_data
+
+
+def config_deserialization(in_config: dict) -> Union[TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig]:
+    """
+    This function change config dictionary to trainable quantizer config.
+    Args:
+        in_config:  A config dictionary of trainable quantizer config.
+
+    Returns: Trainable quantizer configuration object - TrainableQuantizerWeightsConfig or TrainableQuantizerActivationConfig
+
+    """
+    in_config = copy.deepcopy(in_config)
+    if in_config[C.IS_WEIGHTS]:
+        return TrainableQuantizerWeightsConfig(weights_quantization_method=QuantizationMethod(in_config[C.WEIGHTS_QUANTIZATION_METHOD]),
+                                               weights_n_bits=in_config[C.WEIGHTS_N_BITS],
+                                               weights_quantization_params=in_config[C.WEIGHTS_QUANTIZATION_PARAMS],
+                                               enable_weights_quantization=in_config[C.ENABLE_WEIGHTS_QUANTIZATION],
+                                               weights_channels_axis=in_config[C.WEIGHTS_CHANNELS_AXIS],
+                                               weights_per_channel_threshold=in_config[C.WEIGHTS_PER_CHANNEL_THRESHOLD],
+                                               min_threshold=in_config[C.MIN_THRESHOLD])
+    elif in_config[C.IS_ACTIVATIONS]:
+        return TrainableQuantizerActivationConfig(activation_quantization_method=QuantizationMethod(in_config[C.ACTIVATION_QUANTIZATION_METHOD]),
+                                                  activation_n_bits=in_config[C.ACTIVATION_N_BITS],
+                                                  activation_quantization_params=in_config[C.ACTIVATION_QUANTIZATION_PARAMS],
+                                                  enable_activation_quantization=in_config[C.ENABLE_ACTIVATION_QUANTIZATION],
+                                                  min_threshold=in_config[C.MIN_THRESHOLD])
+    else:
+        raise NotImplemented  # pragma: no cover

model_compression_toolkit/quantizers_infrastructure/trainable_infrastructure/keras/quantizer_utils.py

@@ -0,0 +1,48 @@
+# 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.
+# ==============================================================================
+
+import numpy as np
+import tensorflow as tf
+
+
+def int_quantization_with_threshold(data: tf.Tensor,
+                                    n_bits: int,
+                                    signed: bool,
+                                    threshold: np.ndarray,
+                                    eps: float) -> tf.Tensor:
+    """
+    Divides data by threshold and quantize it to integers in the quantization range (depends on signed value).
+
+    Args:
+        data: tensor data.
+        n_bits: number of bits that determines the quantization range.
+        signed: Whether the quantization is signed or not.
+        threshold: threshold for quantization.
+        eps: Small value for numerical stability in division.
+
+    Returns:
+        Uniform Quantized tensor.
+
+    """
+
+    if signed:
+        clip_max = 2 ** (n_bits - 1) - 1
+        clip_min = -2 ** (n_bits - 1)
+    else:
+        clip_max = 2 ** n_bits - 1
+        clip_min = 0
+
+    return tf.clip_by_value((data / (threshold + eps)) * (2 ** (n_bits - int(signed))),
+                            clip_value_max=clip_max, clip_value_min=clip_min)

model_compression_toolkit/quantizers_infrastructure/trainable_infrastructure/pytorch/__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/quantizers_infrastructure/trainable_infrastructure/pytorch/base_pytorch_quantizer.py

@@ -0,0 +1,66 @@
+# Copyright 2022 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, List
+
+from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.core.common.constants import FOUND_TORCH
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import BaseTrainableQuantizer, VAR, GROUP
+from model_compression_toolkit.quantizers_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)
+
+
+        def get_trainable_variables(self, group: VariableGroup) -> List[torch.Tensor]:
+            """
+            Get trainable parameters with specific group from quantizer
+
+            Args:
+                group: Enum of variable group
+
+            Returns:
+                List of trainable variables
+            """
+            quantizer_trainable = []
+            for name, parameter_dict in self.quantizer_parameters.items():
+                quantizer_parameter, parameter_group = parameter_dict[VAR], parameter_dict[GROUP]
+                if quantizer_parameter.requires_grad and parameter_group == group:
+                    quantizer_trainable.append(quantizer_parameter)
+            return quantizer_trainable
+
+else:
+    class BasePytorchTrainableQuantizer(BaseTrainableQuantizer):
+        def __init__(self,
+                     quantization_config: Union[TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig]):
+            super().__init__(quantization_config)
+            Logger.critical('Installing Pytorch is mandatory '
+                            'when using BasePytorchTrainableQuantizer. '
+                            'Could not find torch package.')  # pragma: no cover