mct-nightly 1.8.0.27022023.post430__py3-none-any.whl → 1.8.0.27032023.post403__py3-none-any.whl

model_compression_toolkit/gptq/pytorch/quantizer/soft_rounding/uniform_soft_quantizer.py

@@ -0,0 +1,196 @@
+# 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 torch
+import torch.nn as nn
+from typing import Dict
+import numpy as np
+
+from model_compression_toolkit import quantizers_infrastructure as qi
+from model_compression_toolkit.core.common.target_platform import QuantizationMethod
+from model_compression_toolkit.gptq.common.gptq_config import RoundingType
+from model_compression_toolkit.gptq.pytorch.quantizer.base_pytorch_gptq_quantizer import \
+    BasePytorchGPTQTrainableQuantizer
+from model_compression_toolkit.core.pytorch.utils import to_torch_tensor, torch_tensor_to_numpy
+from model_compression_toolkit.gptq.pytorch.quantizer import quant_utils as qutils
+from model_compression_toolkit.gptq.common.gptq_constants import SOFT_ROUNDING_GAMMA, SOFT_ROUNDING_ZETA, AUXVAR
+from model_compression_toolkit.gptq.pytorch.quantizer.quant_utils import fix_range_to_include_zero
+from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.base_inferable_quantizer import \
+    mark_quantizer
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import \
+    VariableGroup
+from model_compression_toolkit.core.common.constants import RANGE_MAX, RANGE_MIN
+from model_compression_toolkit.qat.common.constants import FQ_MIN, FQ_MAX
+
+def soft_rounding_unifrom_quantizer(input_tensor: torch.Tensor,
+                                    auxvar_tensor: torch.Tensor,
+                                    min_range: torch.Tensor,
+                                    max_range: torch.Tensor,
+                                    num_bits: int) -> torch.Tensor:
+    """
+    Quantize a tensor uniformly for GPTQ quantizers.
+
+    Args:
+        input_tensor: Tensor to quantize. values of this tensor are not changed during gptq.
+        auxvar_tensor: Tensor that manifests the bit shift of the quantized weights due to gptq training.
+        min_range: Tensor with min values to compute the delta grid.
+        max_range: Tensor with max values to compute the delta grid.
+        num_bits: Num of bits to use.
+
+    Returns:
+        A quantized tensor.
+    """
+    # adjusts the quantization range so the quantization grid includes zero.
+    min_range, max_range = fix_range_to_include_zero(min_range, max_range, num_bits)
+    delta = qutils.calculate_delta_uniform(max_range, min_range, num_bits)
+    with torch.no_grad():
+        input_tensor_int = torch.floor(input_tensor / delta)
+    tensor_q = input_tensor_int + auxvar_tensor
+    return delta * qutils.ste_clip(tensor_q,
+                                   min_val=0,
+                                   max_val=2 ** num_bits - 1)
+
+
+@mark_quantizer(quantization_target=qi.QuantizationTarget.Weights,
+                quantization_method=[QuantizationMethod.UNIFORM],
+                quantizer_type=RoundingType.SoftQuantizer)
+class UniformSoftRoundingGPTQ(BasePytorchGPTQTrainableQuantizer):
+    """
+    Trainable uniform quantizer to optimize the rounding of the quantized values using a soft quantization method.
+    """
+
+    def __init__(self,
+                 quantization_config: TrainableQuantizerWeightsConfig,
+                 quantization_parameter_learning: bool = False):
+        """
+        Construct a Pytorch model that utilize a fake weight quantizer of soft-quantizer for symmetric quantizer.
+
+        Args:
+            quantization_config: Trainable weights quantizer config.
+            quantization_parameter_learning (Bool): Whether to learn the min/max ranges or not
+        """
+
+        super().__init__(quantization_config)
+        self.num_bits = quantization_config.weights_n_bits
+        self.per_channel = quantization_config.weights_per_channel_threshold
+
+        self.min_values = quantization_config.weights_quantization_params[RANGE_MIN]
+        self.max_values = quantization_config.weights_quantization_params[RANGE_MAX]
+
+        self.quantization_axis = quantization_config.weights_channels_axis
+        self.quantization_parameter_learning = quantization_parameter_learning
+
+        # gamma and zeta are stretch parameters for computing the rectified sigmoid function.
+        # See: https://arxiv.org/pdf/2004.10568.pdf
+        self.gamma = SOFT_ROUNDING_GAMMA
+        self.zeta = SOFT_ROUNDING_ZETA
+
+    def initialize_quantization(self,
+                                tensor_shape: torch.Size,
+                                name: str,
+                                layer: qi.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.
+        if self.per_channel:
+            min_values = to_torch_tensor(self.min_values)
+            max_values = to_torch_tensor(self.max_values)
+        else:
+            min_values = torch.tensor(self.min_values)
+            max_values = torch.tensor(self.max_values)
+
+        layer.register_parameter(name+"_"+FQ_MIN, nn.Parameter(min_values, requires_grad=self.quantization_parameter_learning))
+        layer.register_parameter(name+"_"+FQ_MAX, nn.Parameter(max_values, requires_grad=self.quantization_parameter_learning))
+
+        w = layer.layer.weight
+        delta = qutils.calculate_delta_uniform(max_values, min_values, self.num_bits)
+        w_clipped_normed = torch.clip(w / delta, 0, 2 ** self.num_bits - 1)
+        rest = w_clipped_normed - torch.floor(w_clipped_normed)  # rest of rounding [0, 1)
+        alpha = -torch.log((self.zeta - self.gamma) / (rest - self.gamma) - 1)  # => sigmoid(alpha) = rest
+        layer.register_parameter(f"{name}_{AUXVAR}", nn.Parameter(alpha, requires_grad=True))
+
+        # Save the quantizer parameters
+        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)
+        self.add_quantizer_variable(AUXVAR, layer.get_parameter(f"{name}_{AUXVAR}"), VariableGroup.WEIGHTS)
+
+
+    def get_soft_targets(self) -> torch.Tensor:
+        """
+        Computes the rectified sigmoid function for the quantization target parameters.
+
+        Returns:
+            A tensor with the soft rounding targets values.
+
+        """
+        scaled_sigmoid = torch.sigmoid(self.get_quantizer_variable(AUXVAR)) * (self.zeta - self.gamma) + self.gamma
+        return torch.clip(scaled_sigmoid, min=0, max=1)
+
+    def get_quant_config(self) -> Dict[str, np.ndarray]:
+        """
+        Returns the config used to edit NodeQuantizationConfig after GPTQ retraining
+
+        Returns:
+            A dictionary of attributes the quantize_config retraining has changed during GPTQ retraining.
+            Keys must match NodeQuantizationConfig attributes
+
+        """
+        min_values = torch_tensor_to_numpy(self.get_quantizer_variable(FQ_MIN))
+        max_values = torch_tensor_to_numpy(self.get_quantizer_variable(FQ_MAX))
+        return {RANGE_MIN:  min_values,
+                RANGE_MAX:  max_values}
+
+    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
+        """
+        auxvar = self.get_quantizer_variable(AUXVAR)
+        min_range = self.get_quantizer_variable(FQ_MIN)
+        max_range = self.get_quantizer_variable(FQ_MAX)
+
+        #####################################################
+        # Soft Rounding
+        #####################################################
+        aux_var = self.get_soft_targets()
+        if not training:
+            aux_var = (aux_var >= 0.5).to(auxvar.dtype)
+
+        #####################################################
+        # Quantized Input
+        #####################################################
+        q_tensor = soft_rounding_unifrom_quantizer(input_tensor=inputs,
+                                                   auxvar_tensor=aux_var,
+                                                   min_range=min_range,
+                                                   max_range=max_range,
+                                                   num_bits=self.num_bits)
+
+
+        return q_tensor

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

@@ -14,7 +14,7 @@
 # ==============================================================================
 import torch
 import torch.nn as nn
-from typing import List, Dict
+from typing import Dict
 import numpy as np
 from model_compression_toolkit.core.common.defaultdict import DefaultDict
 
@@ -30,7 +30,7 @@ from model_compression_toolkit.core.common.constants import THRESHOLD
 from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.base_inferable_quantizer import \
     mark_quantizer
-
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.quant_utils import \
     get_threshold_reshape_shape
 
@@ -104,23 +104,19 @@ class STEWeightGPTQQuantizer(BasePytorchGPTQTrainableQuantizer):
         self.quantization_axis = quantization_config.weights_channels_axis
         self.power_of_two = quantization_config.weights_quantization_method == QuantizationMethod.POWER_OF_TWO
         self.max_lsbs_change = max_lsbs_change_map.get(self.num_bits)
-        self.quantizer_parameters = {}
 
 
     def initialize_quantization(self,
                                 tensor_shape: torch.Size,
                                 name: str,
-                                layer: qi.PytorchQuantizationWrapper) -> Dict[str, nn.Parameter]:
+                                layer: qi.PytorchQuantizationWrapper):
         """
-        Return a dictionary of quantizer parameters and their names.
+        Add quantizer parameters to the quantizer parameters dictionary
 
         Args:
             tensor_shape: tensor shape of the quantized tensor.
             name: Tensor name.
             layer: Layer to quantize.
-
-        Returns:
-            Dictionary of parameters names to the variables.
         """
 
         layer.register_parameter(f"{name}_{PTQ_THRESHOLD}",
@@ -131,27 +127,9 @@ class STEWeightGPTQQuantizer(BasePytorchGPTQTrainableQuantizer):
                                                                   requires_grad=True))
 
         # save the quantizer added parameters for later calculations
-        self.quantizer_parameters = {PTQ_THRESHOLD: layer.get_parameter(f"{name}_{PTQ_THRESHOLD}"),
-                                     AUXVAR: layer.get_parameter(f"{name}_{AUXVAR}")}
-
-        return self.quantizer_parameters
-
-
-    def get_aux_variable(self) -> List[torch.Tensor]:
-        """
-        This function return a list with the quantizer's quantization auxiliary variables.
-
-        Returns: A list with the quantization auxiliary variables.
-        """
-        return [self.quantizer_parameters.get(AUXVAR)]
+        self.add_quantizer_variable(PTQ_THRESHOLD, layer.get_parameter(f"{name}_{PTQ_THRESHOLD}"), VariableGroup.QPARAMS)
+        self.add_quantizer_variable(AUXVAR, layer.get_parameter(f"{name}_{AUXVAR}"), VariableGroup.WEIGHTS)
 
-    def get_quantization_variable(self) -> List[torch.Tensor]:
-        """
-        This function return a list with the quantizer's quantization parameters variables.
-
-        Returns: A list with the quantization parameters.
-        """
-        return [self.quantizer_parameters.get(PTQ_THRESHOLD)]
 
     def get_quant_config(self) -> Dict[str, np.ndarray]:
         """
@@ -162,7 +140,7 @@ class STEWeightGPTQQuantizer(BasePytorchGPTQTrainableQuantizer):
             Keys must match NodeQuantizationConfig attributes
 
         """
-        old_threshold = self.quantizer_parameters[PTQ_THRESHOLD]
+        old_threshold = self.get_quantizer_variable(PTQ_THRESHOLD)
         return {THRESHOLD: torch_tensor_to_numpy(old_threshold).reshape(self.threshold_shape)}
 
     def __call__(self,
@@ -178,8 +156,8 @@ class STEWeightGPTQQuantizer(BasePytorchGPTQTrainableQuantizer):
         Returns:
             quantized tensor
         """
-        auxvar = self.quantizer_parameters[AUXVAR]
-        ptq_threshold_tensor = self.quantizer_parameters[PTQ_THRESHOLD]
+        auxvar = self.get_quantizer_variable(AUXVAR)
+        ptq_threshold_tensor = self.get_quantizer_variable(PTQ_THRESHOLD)
 
         if self.per_channel:
             reshape_shape = get_threshold_reshape_shape(inputs.shape,

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

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Dict, Union
+from typing import Union
 
 import numpy as np
 import tensorflow as tf
@@ -32,6 +32,7 @@ from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructur
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers import \
     WeightsPOTInferableQuantizer, WeightsSymmetricInferableQuantizer, ActivationPOTInferableQuantizer, \
     ActivationSymmetricInferableQuantizer
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
 @mark_quantizer(quantization_target=qi.QuantizationTarget.Weights,
@@ -76,22 +77,19 @@ class STEWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
         max_int = (2 ** (self.num_bits - int(C.WEIGHTS_SIGNED))) - 1
         self.min = delta * min_int
         self.max = delta * max_int
-        self.quantizer_parameters = {}
+
 
     def initialize_quantization(self,
                                 tensor_shape: TensorShape,
                                 name: str,
-                                layer: qi.KerasQuantizationWrapper) -> Dict[str, tf.Variable]:
+                                layer: qi.KerasQuantizationWrapper):
         """
-        Add min and max variables to layer.
-        Args:
-            tensor_shape: Tensor shape the quantizer quantize.
-            name: Prefix of variables names.
-            layer: Layer to add the variables to. The variables are saved
-            in the layer's scope.
+        Add quantizer parameters to the quantizer parameters dictionary
 
-        Returns:
-            Dictionary of new variables.
+        Args:
+            tensor_shape: tensor shape of the quantized tensor.
+            name: Tensor name.
+            layer: Layer to quantize.
         """
         ptq_threshold_tensor = layer.add_weight(
             name + THRESHOLD_TENSOR,
@@ -115,9 +113,9 @@ class STEWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
         fq_max.assign(self.max)
 
         # save the quantizer added parameters for later calculations
-        self.quantizer_parameters = {THRESHOLD_TENSOR: ptq_threshold_tensor,
-                                     FQ_MIN: fq_min, FQ_MAX: fq_max}
-        return self.quantizer_parameters
+        self.add_quantizer_variable(THRESHOLD_TENSOR, ptq_threshold_tensor, VariableGroup.QPARAMS)
+        self.add_quantizer_variable(FQ_MIN, fq_min, VariableGroup.QPARAMS)
+        self.add_quantizer_variable(FQ_MAX, fq_max, VariableGroup.QPARAMS)
 
     def __call__(self,
                  inputs: tf.Tensor,
@@ -134,8 +132,8 @@ class STEWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
             The quantized tensor.
         """
 
-        _min = self.quantizer_parameters[FQ_MIN]
-        _max = self.quantizer_parameters[FQ_MAX]
+        _min = self.get_quantizer_variable(FQ_MIN)
+        _max = self.get_quantizer_variable(FQ_MAX)
         if self.channel_axis:
             if self.perm_vec:
                 inputs = tf.transpose(inputs, perm=self.perm_vec)
@@ -157,7 +155,7 @@ class STEWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
             BaseKerasInferableQuantizer object.
         """
         if self.power_of_two:
-            pot_threshold = 2 ** np.ceil(np.log2(self.quantizer_parameters[THRESHOLD_TENSOR]))
+            pot_threshold = 2 ** np.ceil(np.log2(self.get_quantizer_variable(THRESHOLD_TENSOR)))
             return WeightsPOTInferableQuantizer(num_bits=self.num_bits,
                                                 threshold=list(pot_threshold.flatten()),
                                                 per_channel=self.per_channel,
@@ -165,8 +163,7 @@ class STEWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
                                                 input_rank=len(self.threshold_shape))
         else:
             return WeightsSymmetricInferableQuantizer(num_bits=self.num_bits,
-                                                      threshold=list(self.quantizer_parameters[
-                                                                         THRESHOLD_TENSOR].numpy().flatten()),
+                                                      threshold=list(self.get_quantizer_variable(THRESHOLD_TENSOR).numpy().flatten()),
                                                       per_channel=self.per_channel,
                                                       channel_axis=self.channel_axis,
                                                       input_rank=len(self.threshold_shape))
@@ -203,22 +200,18 @@ class STEActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
         max_int = (2 ** (self.num_bits - int(self.signed))) - 1
         self.min = delta * min_int
         self.max = delta * max_int
-        self.quantizer_parameters = {}
 
     def initialize_quantization(self,
                                 tensor_shape: TensorShape,
                                 name: str,
-                                layer: qi.KerasQuantizationWrapper) -> Dict[str, tf.Variable]:
+                                layer: qi.KerasQuantizationWrapper):
         """
-        Add min and max variables to layer.
-        Args:
-            tensor_shape: Tensor shape the quantizer quantize.
-            name: Prefix of variables names.
-            layer: Layer to add the variables to. The variables are saved
-            in the layer's scope.
+        Add quantizer parameters to the quantizer parameters dictionary
 
-        Returns:
-            Dictionary of new variables.
+        Args:
+            tensor_shape: tensor shape of the quantized tensor.
+            name: Tensor name.
+            layer: Layer to quantize.
         """
         ptq_threshold_tensor = layer.add_weight(
             name + THRESHOLD_TENSOR,
@@ -242,9 +235,10 @@ class STEActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
         fq_max.assign(self.max)
 
         # save the quantizer added parameters for later calculations
-        self.quantizer_parameters = {THRESHOLD_TENSOR: ptq_threshold_tensor,
-                                     FQ_MIN: fq_min, FQ_MAX: fq_max}
-        return self.quantizer_parameters
+        self.add_quantizer_variable(THRESHOLD_TENSOR, ptq_threshold_tensor, VariableGroup.QPARAMS)
+        self.add_quantizer_variable(FQ_MIN, fq_min, VariableGroup.QPARAMS)
+        self.add_quantizer_variable(FQ_MAX, fq_max, VariableGroup.QPARAMS)
+
 
     def __call__(self,
                  inputs: tf.Tensor,
@@ -259,8 +253,8 @@ class STEActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
             The quantized tensor.
         """
 
-        _min = self.quantizer_parameters[FQ_MIN]
-        _max = self.quantizer_parameters[FQ_MAX]
+        _min = self.get_quantizer_variable(FQ_MIN)
+        _max = self.get_quantizer_variable(FQ_MAX)
         q_tensor = tf.quantization.fake_quant_with_min_max_vars(inputs, _min, _max,
                                                                 num_bits=self.num_bits)
 
@@ -275,7 +269,7 @@ class STEActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
         """
 
         if self.power_of_two:
-            pot_threshold = 2 ** np.ceil(np.log2(self.quantizer_parameters[THRESHOLD_TENSOR]))
+            pot_threshold = 2 ** np.ceil(np.log2(self.get_quantizer_variable(THRESHOLD_TENSOR)))
             return ActivationPOTInferableQuantizer(num_bits=self.num_bits,
                                                       # In activation quantization is per-tensor only - thus we pass
                                                       # the threshold as a list with a len of 1
@@ -285,6 +279,5 @@ class STEActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
             return ActivationSymmetricInferableQuantizer(num_bits=self.num_bits,
                                                          # In activation quantization is per-tensor only - thus we
                                                          # pass the threshold as a list with a len of 1
-                                                         threshold=[
-                                                             self.quantizer_parameters[THRESHOLD_TENSOR].numpy()],
+                                                         threshold=[self.get_quantizer_variable(THRESHOLD_TENSOR).numpy()],
                                                          signed=self.signed)

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

@@ -12,9 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
-from typing import Dict
-
 import numpy as np
 import tensorflow as tf
 from tensorflow.python.framework.tensor_shape import TensorShape
@@ -32,6 +29,7 @@ from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructur
     mark_quantizer
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers import \
     BaseKerasInferableQuantizer, WeightsUniformInferableQuantizer, ActivationUniformInferableQuantizer
+from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
 
 @mark_quantizer(quantization_target=qi.QuantizationTarget.Weights,
@@ -70,22 +68,17 @@ class STEUniformWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
         else:
             self.perm_vec = None
 
-        self.quantizer_parameters = {}
-
     def initialize_quantization(self,
                                 tensor_shape: TensorShape,
                                 name: str,
-                                layer: qi.KerasQuantizationWrapper) -> Dict[str, tf.Variable]:
+                                layer: qi.KerasQuantizationWrapper):
         """
-        Add min and max variables to layer.
-        Args:
-            tensor_shape: Tensor shape the quantizer quantize.
-            name: Prefix of variables names.
-            layer: Layer to add the variables to. The variables are saved
-            in the layer's scope.
+        Add quantizer parameters to the quantizer parameters dictionary
 
-        Returns:
-            Dictionary of new variables.
+        Args:
+            tensor_shape: tensor shape of the quantized tensor.
+            name: Tensor name.
+            layer: Layer to quantize.
         """
         fq_min = layer.add_weight(
             name + FQ_MIN,
@@ -102,8 +95,9 @@ class STEUniformWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
         fq_max.assign(self.max)
 
         # save the quantizer added parameters for later calculations
-        self.quantizer_parameters = {FQ_MIN: fq_min, FQ_MAX: fq_max}
-        return self.quantizer_parameters
+        self.add_quantizer_variable(FQ_MIN, fq_min, VariableGroup.QPARAMS)
+        self.add_quantizer_variable(FQ_MAX, fq_max, VariableGroup.QPARAMS)
+
 
     def __call__(self, inputs: tf.Tensor,
                  training: bool):
@@ -117,8 +111,8 @@ class STEUniformWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
             The quantized tensor.
         """
 
-        _min = self.quantizer_parameters[FQ_MIN]
-        _max = self.quantizer_parameters[FQ_MAX]
+        _min = self.get_quantizer_variable(FQ_MIN)
+        _max = self.get_quantizer_variable(FQ_MAX)
         _min, _max = adjust_range_to_include_zero(_min, _max, self.num_bits)
 
         if self.per_channel:
@@ -142,8 +136,8 @@ class STEUniformWeightQATQuantizer(BaseKerasQATTrainableQuantizer):
         Returns:
             BaseKerasInferableQuantizer object.
         """
-        min_range, max_range = fix_range_to_include_zero(self.quantizer_parameters[FQ_MIN].numpy(),
-                                                         self.quantizer_parameters[FQ_MAX].numpy(),
+        min_range, max_range = fix_range_to_include_zero(self.get_quantizer_variable(FQ_MIN).numpy(),
+                                                         self.get_quantizer_variable(FQ_MAX).numpy(),
                                                          self.num_bits)
         return WeightsUniformInferableQuantizer(num_bits=self.num_bits,
                                                 min_range=list(min_range.flatten()),
@@ -174,22 +168,18 @@ class STEUniformActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
         self.num_bits = quantization_config.activation_n_bits
         self.min_range = quantization_config.activation_quantization_params[C.RANGE_MIN]
         self.max_range = quantization_config.activation_quantization_params[C.RANGE_MAX]
-        self.quantizer_parameters = {}
 
     def initialize_quantization(self,
                                 tensor_shape: TensorShape,
                                 name: str,
-                                layer: qi.KerasQuantizationWrapper) -> Dict[str, tf.Variable]:
+                                layer: qi.KerasQuantizationWrapper):
         """
-        Add min and max variables to layer.
-        Args:
-            tensor_shape: Tensor shape the quantizer quantize.
-            name: Prefix of variables names.
-            layer: Layer to add the variables to. The variables are saved
-            in the layer's scope.
+        Add quantizer parameters to the quantizer parameters dictionary
 
-        Returns:
-            Dictionary of new variables.
+        Args:
+            tensor_shape: tensor shape of the quantized tensor.
+            name: Tensor name.
+            layer: Layer to quantize.
         """
         fq_min = layer.add_weight(
             name + FQ_MIN,
@@ -206,8 +196,9 @@ class STEUniformActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
         fq_max.assign(self.max_range)
 
         # save the quantizer added parameters for later calculations
-        self.quantizer_parameters = {FQ_MIN: fq_min, FQ_MAX: fq_max}
-        return self.quantizer_parameters
+        self.add_quantizer_variable(FQ_MIN, fq_min, VariableGroup.QPARAMS)
+        self.add_quantizer_variable(FQ_MAX, fq_max, VariableGroup.QPARAMS)
+
 
     def __call__(self,
                  inputs: tf.Tensor,
@@ -222,8 +213,8 @@ class STEUniformActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
             The quantized tensor.
         """
 
-        _min = self.quantizer_parameters[FQ_MIN]
-        _max = self.quantizer_parameters[FQ_MAX]
+        _min = self.get_quantizer_variable(FQ_MIN)
+        _max = self.get_quantizer_variable(FQ_MAX)
         _min, _max = adjust_range_to_include_zero(_min, _max, self.num_bits)
         q_tensor = tf.quantization.fake_quant_with_min_max_vars(inputs, _min, _max,
                                                                 num_bits=self.num_bits)
@@ -237,8 +228,8 @@ class STEUniformActivationQATQuantizer(BaseKerasQATTrainableQuantizer):
         Returns:
             BaseKerasInferableQuantizer object.
         """
-        min_range, max_range = fix_range_to_include_zero(self.quantizer_parameters[FQ_MIN].numpy(),
-                                                         self.quantizer_parameters[FQ_MAX].numpy(),
+        min_range, max_range = fix_range_to_include_zero(self.get_quantizer_variable(FQ_MIN).numpy(),
+                                                         self.get_quantizer_variable(FQ_MAX).numpy(),
                                                          self.num_bits)
         return ActivationUniformInferableQuantizer(num_bits=self.num_bits,
                                                    # In activation quantization is per-tensor only - thus we pass