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

model_compression_toolkit/gptq/pytorch/quantizer/gumbel_rounding/base_gumbel_weights_quantizer.py

@@ -1,157 +0,0 @@
-# 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 abc import abstractmethod
-import torch
-import numpy as np
-from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfigV2
-from model_compression_toolkit.core.common.logger import Logger
-from model_compression_toolkit.gptq.pytorch.quantizer.gptq_quantizer import BaseWeightQuantizer
-from model_compression_toolkit.core.common.quantization.node_quantization_config import NodeWeightsQuantizationConfig
-from model_compression_toolkit.gptq.pytorch.quantizer.quant_utils import sample_gumbel
-from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
-from model_compression_toolkit.core.common.target_platform.op_quantization_config import QuantizationMethod
-from model_compression_toolkit.gptq.pytorch.quantizer.quant_utils import ste_clip
-
-P_INIT = 0.01
-GR_SHIFT_BASE = 2
-
-
-def init_aux_var(ceil_indicator: torch.Tensor, w_shape: torch.Size, m: int, p: float = P_INIT) -> torch.Tensor:
-    """
-    This function generate a random pi matrix for Gumbel Rounding
-    Args:
-        ceil_indicator: An array of indicator if the value should be ceil or floor.
-        w_shape(torch.Size): A list of integers that represent the shape of the weights tensor to be quantization
-        m(int):  An integer that define the number of shift
-        p(float): A floating point number that represent the probability of non-round options of pi matrix.
-
-    Returns: A torch tensor of pi tensor
-
-    """
-
-    if m < 2:
-        Logger.error("m must be larger than two")
-    if m % 2 != 0:
-        Logger.error("m must be module two")
-    m_hat = m // 2 - 1
-    shift = -np.log(-np.log(1 - p))
-    n = np.random.randn(*[m, *w_shape]) * np.sqrt(np.power(np.pi, 2) / 6)
-    n = n.reshape([m, -1]).T
-    ceil_indicator = ceil_indicator.cpu().numpy().flatten()
-    n[np.arange(ceil_indicator.size), ceil_indicator + m_hat] += shift
-    n = n.T.reshape(*[m, *w_shape])
-    return torch.from_numpy(n).float()
-
-
-def init_shift_var(m: int) -> torch.Tensor:
-    """
-    This function generate a tensor of 2*m+1 from -m to m
-    Args:
-        m: An integer value the represent m
-
-    Returns: A tensor of size m
-
-    """
-    m_hat = m // 2
-    aux_index_shift = [-m_hat + i + 1 for i in range(m)]
-    return torch.Tensor(aux_index_shift)
-
-
-class BaseGumbelWeightQuantizer(BaseWeightQuantizer):
-    """
-    Base class that implements a quantizer with trainable parameters to be used for GPTQ training.
-    """
-
-    def __init__(self,
-                 weights_quantization_cfg: NodeWeightsQuantizationConfig,
-                 gptq_config: GradientPTQConfigV2,
-                 weight_shape: torch.Size):
-        """
-        Construct a Pytorch model that utilize a fake weight quantizer of Gumbel rounding
-        Args:
-            weights_quantization_cfg: Configuration of weight quantization
-            gptq_config: GradientPTQConfigV2 object with parameters about the tuning process.
-            weight_shape: weight shape for auxiliary tensor creation.
-        """
-        super().__init__()
-        self.power_of_two = QuantizationMethod.POWER_OF_TWO == weights_quantization_cfg.weights_quantization_method
-        self.reshape_aux_shift = [-1, *[1 for _ in range(len(weight_shape))]]
-        self.num_bits = weights_quantization_cfg.weights_n_bits
-        self.weight_shape = weight_shape
-        self.max_delta_change = gptq_config.lsb_change_per_bit_width.get(self.num_bits)
-        self.quantization_parameter_learning = gptq_config.quantization_parameters_learning
-        self.m = GR_SHIFT_BASE * self.max_delta_change + GR_SHIFT_BASE
-        self.minimal_temp = gptq_config.quantizer_config.minimal_temp
-        self.maximal_temp = gptq_config.quantizer_config.maximal_temp
-        self.temperature_learning = gptq_config.quantizer_config.temperature_learning
-        self.cycle_iterations = max(1, int(gptq_config.n_epochs / gptq_config.quantizer_config.n_cycles))
-        self.shift_tensor = to_torch_tensor(init_shift_var(self.m))
-        self.tau = None
-        self.g_t = 0
-        self.p_t = None
-        self.n_iter = 0
-        self.update_gumbel_param = True
-        scale = self.cycle_iterations / (-2 * np.log(0.001))
-
-        self.gumbel_scale = gptq_config.quantizer_config.gumbel_scale
-        self.gumbel_scale_per_bitwidth = gptq_config.quantizer_config.gumbel_scale_per_bitwidth
-
-        def tau_function(i: int) -> float:
-            """
-            A function that generates the gumbel temperature.
-            Args:
-                i: An int that represents the current iteration number
-
-            Returns: A temperature value.
-
-            """
-            if i < (self.cycle_iterations - 1):
-                index = ((i + 1) % self.cycle_iterations) / scale
-            else:
-                index = (i % self.cycle_iterations) / scale
-
-            x = np.exp(-index)
-            return self.minimal_temp + (self.maximal_temp - self.minimal_temp) * x
-
-        self.tau_function = tau_function
-
-    def get_gumbel_probability(self) -> torch.Tensor:
-        """
-        A function that return the gumbel probability value.
-        Returns: gumbel probability
-        """
-        return self.p_t
-
-    def update_iteration(self, training):
-        """
-        A function that update parameters for GPTQ fine-tuning
-        Args:
-            training: whether in training mode or not
-        """
-        if self.temperature_learning:
-            self.tau = ste_clip(self.temp_tensor, self.minimal_temp, self.maximal_temp)
-        else:
-            self.tau = self.tau_function(self.n_iter)
-        if self.update_gumbel_param and training:
-            self.n_iter += 1
-            self.g_t = sample_gumbel([self.m, *self.weight_shape])
-
-    @abstractmethod
-    def get_temperature_variable(self) -> Union[torch.Tensor, List]:
-        """
-        Returns temperature trainable variables
-        """
-        raise Logger.error(f"{self.__class__.__name__} have to implement this abstract function.")

model_compression_toolkit/gptq/pytorch/quantizer/gumbel_rounding/sym_gumbel_weights_quantizer.py

@@ -1,150 +0,0 @@
-# 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 torch
-import torch.nn as nn
-from typing import List, Union
-from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig
-from model_compression_toolkit.gptq.pytorch.quantizer.gumbel_rounding.base_gumbel_weights_quantizer import BaseGumbelWeightQuantizer, init_aux_var
-from model_compression_toolkit.core.pytorch.utils import to_torch_tensor, torch_tensor_to_numpy
-from model_compression_toolkit.gptq.pytorch.quantizer.quant_utils import symmetric_quantizer
-from model_compression_toolkit.gptq.pytorch.quantizer.quant_utils import ste_clip, ste_gumbel, gumbel_softmax, power_of_two_max
-from model_compression_toolkit.gptq.common.gptq_constants import AUXVAR, THRESHOLD_TENSOR, TEMP, SCALE_TENSOR
-from model_compression_toolkit.core.common.quantization.node_quantization_config import NodeWeightsQuantizationConfig
-from model_compression_toolkit.core.common.constants import THRESHOLD
-
-
-class SymmetricGumbelWeightQuantizer(BaseGumbelWeightQuantizer):
-    """
-    Class that implements a quantizer with trainable parameters to be used for GPTQ training.
-    """
-
-    def __init__(self,
-                 weights_quantization_cfg: NodeWeightsQuantizationConfig,
-                 gptq_config: GradientPTQConfig,
-                 weight: torch.Tensor):
-        """
-        Construct a Pytorch model that utilize a fake weight quantizer of Symmetric Gumbel rounding
-        Args:
-            weights_quantization_cfg: Configuration of weight quantization
-            gptq_config: GradientPTQConfig object with parameters about the tuning process.
-            weight: weight for auxiliary tensor creation.
-        """
-        super().__init__(weights_quantization_cfg, gptq_config, weight.shape)
-        self.signed = True
-        self.min_int = -int(self.signed) * (2 ** (self.num_bits - int(self.signed)))
-        self.max_int = (2 ** (self.num_bits - int(self.signed))) - 1
-        self.threshold_tensor = to_torch_tensor(weights_quantization_cfg.weights_quantization_params.get(THRESHOLD))
-        self.scale_tensor = torch.ones(self.weight_shape)
-
-        # Set trainable tensors
-        self.set_trainable_params(weight)
-
-
-    def set_trainable_params(self, weight: torch.nn.Parameter):
-        """
-        A function to set a list of trainable parameters of the quantizer for GPTQ retraining
-        Args:
-            weight: weight for auxiliary tensor creation.
-        """
-        q_error = weight - symmetric_quantizer(weight,
-                                               self.threshold_tensor,
-                                               num_bits=self.num_bits,
-                                               signed=True,
-                                               power_of_two=self.power_of_two)
-        ceil_indicator = (q_error < 0).int()  # Negative error means the choosen point is rounded to ceil.
-        self.aux_tensor = nn.Parameter(to_torch_tensor(init_aux_var(ceil_indicator, self.weight_shape, self.m)), requires_grad=True)
-        self.trainable_params.update({AUXVAR: self.aux_tensor})
-        self.temp_tensor = nn.Parameter(to_torch_tensor(self.maximal_temp*torch.ones([1,*self.weight_shape])), requires_grad=True)
-        self.trainable_params.update({TEMP: self.temp_tensor})
-        if self.quantization_parameter_learning and not self.power_of_two:
-            self.scale_tensor = nn.Parameter(to_torch_tensor(self.scale_tensor), requires_grad=True)
-            self.trainable_params.update({SCALE_TENSOR: self.scale_tensor})
-        elif self.quantization_parameter_learning:
-            self.threshold_tensor = nn.Parameter(self.threshold_tensor, requires_grad=True)
-            self.trainable_params.update({THRESHOLD_TENSOR: self.threshold_tensor})
-        else:
-            self.trainable_params.update({THRESHOLD_TENSOR: self.threshold_tensor})
-
-    def get_aux_variable(self) -> torch.Tensor:
-        """
-        Returns auxiliary trainable variables
-        """
-        return self.trainable_params.get(AUXVAR)
-
-    def get_quantization_variable(self) -> Union[torch.Tensor, List]:
-        """
-        Returns quantization trainable variables
-        """
-        if self.quantization_parameter_learning and not self.power_of_two:
-            return [self.trainable_params.get(SCALE_TENSOR)]
-        else:
-            return [self.trainable_params.get(THRESHOLD_TENSOR)]
-
-    def get_temperature_variable(self) -> Union[torch.Tensor, List]:
-        """
-        Returns temperature trainable variables
-        """
-        return self.trainable_params.get(TEMP)
-
-    def get_weight_quant_params(self) -> dict:
-        """
-        Returns weight quantization dictionary params
-        """
-        threshold_tensor = self.threshold_tensor
-        if self.power_of_two:
-            threshold_tensor = power_of_two_max(threshold_tensor)
-        elif self.quantization_parameter_learning:
-            threshold_tensor = threshold_tensor*self.scale_tensor
-        return {THRESHOLD: torch_tensor_to_numpy(threshold_tensor.detach())}
-
-    def forward(self, w: nn.Parameter, training:bool = True) -> nn.Parameter:
-        """
-        Weight fake quantizer
-        Args:
-            w: weights to quantize.
-            training: whether in training mode or not
-        Returns:
-            quantized weights
-        """
-        self.update_iteration(training)
-
-        #####################################################
-        # Gumbel Softmax
-        #####################################################
-        if training:
-            gumbel_scale = self.gumbel_scale if self.gumbel_scale_per_bitwidth is None \
-                else self.gumbel_scale_per_bitwidth.get(self.num_bits, self.gumbel_scale)
-            self.p_t = gumbel_softmax(self.aux_tensor, self.tau, self.g_t, gumbel_scale=gumbel_scale)
-        else:
-            self.p_t = ste_gumbel(gumbel_softmax(self.aux_tensor, self.minimal_temp, 0))
-
-        auxhat_tensor = torch.sum(self.p_t * self.shift_tensor.reshape(self.reshape_aux_shift), dim=0)
-
-        #####################################################
-        # Quantizer
-        #####################################################
-        threshold_tensor = self.threshold_tensor
-        if self.power_of_two:
-            threshold_tensor = power_of_two_max(threshold_tensor)
-        delta_tensor = threshold_tensor / (2 ** (self.num_bits-int(self.signed)))
-        w0 = torch.floor(w / delta_tensor).detach()
-        w1 = w0 + auxhat_tensor
-        w2 = ste_clip(w1, min_val=self.min_int, max_val=self.max_int)
-        w_q = delta_tensor * w2
-        # Scale
-        if self.quantization_parameter_learning and not self.power_of_two:
-            w_q *= self.scale_tensor
-        return w_q
-

model_compression_toolkit/gptq/pytorch/quantizer/gumbel_rounding/uniform_gumbel_weights_quantizer.py

@@ -1,143 +0,0 @@
-# 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 numpy as np
-import torch
-import torch.nn as nn
-from typing import List, Union, Dict
-from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig
-from model_compression_toolkit.gptq.pytorch.quantizer.gumbel_rounding.base_gumbel_weights_quantizer import BaseGumbelWeightQuantizer, init_aux_var
-from model_compression_toolkit.core.pytorch.utils import to_torch_tensor, torch_tensor_to_numpy
-from model_compression_toolkit.gptq.pytorch.quantizer.quant_utils import uniform_quantizer, fix_range_to_include_zero
-from model_compression_toolkit.gptq.pytorch.quantizer.quant_utils import ste_clip, ste_gumbel, gumbel_softmax
-from model_compression_toolkit.gptq.common.gptq_constants import AUXVAR, PTQ_MAX_RANGE, PTQ_MIN_RANGE, TEMP
-from model_compression_toolkit.core.common.quantization.node_quantization_config import NodeWeightsQuantizationConfig
-from model_compression_toolkit.core.common.constants import RANGE_MAX, RANGE_MIN
-
-class UniformGumbelWeightQuantizer(BaseGumbelWeightQuantizer):
-    """
-    Class that implements a quantizer with trainable parameters to be used for GPTQ training.
-    """
-
-    def __init__(self,
-                 weights_quantization_cfg: NodeWeightsQuantizationConfig,
-                 gptq_config: GradientPTQConfig,
-                 weight: torch.nn.Parameter):
-        """
-        Construct a Pytorch model that utilize a fake weight quantizer of Uniform Gumbel rounding
-        Args:
-            weights_quantization_cfg: Configuration of weight quantization
-            gptq_config: GradientPTQConfig object with parameters about the tuning process.
-            weight: weight for auxiliary tensor creation.
-        """
-        super().__init__(weights_quantization_cfg, gptq_config, weight.shape)
-        self.min_int = 0
-        self.max_int = 2**self.num_bits - 1
-        self.max_range_tensor = weights_quantization_cfg.weights_quantization_params.get(RANGE_MAX)
-        self.min_range_tensor = weights_quantization_cfg.weights_quantization_params.get(RANGE_MIN)
-
-        # Set trainable tensors
-        self.set_trainable_params(weight)
-
-
-    def set_trainable_params(self, weight: torch.nn.Parameter):
-        """
-        A function to set a list of trainable parameters of the quantizer for GPTQ retraining
-        """
-        self.temp_tensor = nn.Parameter(to_torch_tensor(self.maximal_temp*torch.ones([1,*self.weight_shape])), requires_grad=True)
-        self.trainable_params.update({TEMP: self.temp_tensor})
-        self.max_range_tensor = nn.Parameter(to_torch_tensor(self.max_range_tensor), requires_grad=self.quantization_parameter_learning)
-        self.trainable_params.update({PTQ_MAX_RANGE: self.max_range_tensor})
-        self.min_range_tensor = nn.Parameter(to_torch_tensor(self.min_range_tensor), requires_grad=self.quantization_parameter_learning)
-        self.trainable_params.update({PTQ_MIN_RANGE: self.min_range_tensor})
-        q_error = weight - uniform_quantizer(weight,
-                                             self.min_range_tensor,
-                                             self.max_range_tensor,
-                                             n_bits=self.num_bits)
-        ceil_indicator = (q_error < 0).int()  # Negative error means the choosen point is rounded to ceil.
-        self.aux_tensor = nn.Parameter(to_torch_tensor(init_aux_var(ceil_indicator, self.weight_shape, self.m)), requires_grad=True)
-        self.trainable_params.update({AUXVAR: self.aux_tensor})
-
-    def get_aux_variable(self) -> torch.Tensor:
-        """
-        Returns auxiliary trainable variables
-        """
-        return self.trainable_params.get(AUXVAR)
-
-    def get_quantization_variable(self) -> Union[torch.Tensor, List]:
-        """
-        Returns quantization trainable variables
-        """
-        return [self.trainable_params.get(PTQ_MAX_RANGE), self.trainable_params.get(PTQ_MIN_RANGE)]
-
-    def get_temperature_variable(self) -> Union[torch.Tensor, List]:
-        """
-        Returns temperature trainable variables
-        """
-        return self.trainable_params.get(TEMP)
-
-    def get_weight_quant_params(self) -> Dict[str, np.ndarray]:
-        """
-        Returns weight quantization dictionary params
-        """
-        max_range_tensor = self.max_range_tensor
-        min_range_tensor = self.min_range_tensor
-        return {PTQ_MAX_RANGE: torch_tensor_to_numpy(max_range_tensor.detach()),
-                PTQ_MIN_RANGE: torch_tensor_to_numpy(min_range_tensor.detach())}
-
-    def forward(self, w: nn.Parameter, training:bool = True) -> nn.Parameter:
-        """
-        Weight fake quantizer
-        Args:
-            w: weights to quantize.
-            training: whether in training mode or not
-        Returns:
-            quantized weights
-        """
-        self.update_iteration(training)
-
-        #####################################################
-        # Gumbel Softmax
-        #####################################################
-        if training:
-            self.p_t = gumbel_softmax(self.aux_tensor, self.tau, self.g_t)
-        else:
-            self.p_t = ste_gumbel(gumbel_softmax(self.aux_tensor, self.minimal_temp, 0))
-
-        auxhat_tensor = torch.sum(self.p_t * self.shift_tensor.reshape(self.reshape_aux_shift), dim=0)
-
-        #####################################################
-        # Quantizer
-        #####################################################
-        max_range_tensor = self.max_range_tensor
-        min_range_tensor = self.min_range_tensor
-
-        # adjusts the quantization rage so the quantization grid include zero.
-        a, b = fix_range_to_include_zero(min_range_tensor, max_range_tensor, self.num_bits)
-
-        # Compute the step size of quantized values.
-        delta_tensor = (b - a) / (2 ** self.num_bits - 1)
-
-        # Apply rounding
-        w0 = torch.floor((w - a) / delta_tensor).detach()  # Apply rounding
-
-        w1 = w0 + auxhat_tensor
-
-        # Clip data in range
-        w2 = ste_clip(w1, min_val=self.min_int, max_val=self.max_int)
-
-        # Quantize the data between min/max of quantization range.
-        w_q = delta_tensor * w2 + a
-        return w_q
-

model_compression_toolkit/gptq/pytorch/quantizer/quantizer_wrapper.py

@@ -1,103 +0,0 @@
-# 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 torch
-import torch.nn as nn
-from model_compression_toolkit.core.common import BaseNode, Logger
-from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig, RoundingType
-from model_compression_toolkit.gptq.pytorch.quantizer.gptq_quantizer import BaseWeightQuantizer
-from model_compression_toolkit.gptq.pytorch.quantizer.ste_rounding.ste_weights_quantizer import STEWeightQuantizer
-from model_compression_toolkit.gptq.pytorch.quantizer.gumbel_rounding.sym_gumbel_weights_quantizer import SymmetricGumbelWeightQuantizer
-from model_compression_toolkit.gptq.pytorch.quantizer.gumbel_rounding.uniform_gumbel_weights_quantizer import UniformGumbelWeightQuantizer
-from model_compression_toolkit.core.pytorch.back2framework.instance_builder import node_builder
-from model_compression_toolkit.core.pytorch.constants import KERNEL
-from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
-from model_compression_toolkit.core.common.target_platform.op_quantization_config import QuantizationMethod
-
-
-class WeightQuantizerWrapper(nn.Module):
-
-    def __init__(self, node: BaseNode, gptq_config: GradientPTQConfig, weight_quantizer: BaseWeightQuantizer):
-        """
-        Construct a Pytorch model that constitutes as a wrapper for a Pytorch layer, built from a given graph node.
-        Args:
-            node: Node to build its Pytorch quantizer wrapper.
-            gptq_config: GradientPTQConfig object with parameters about the tuning process.
-            weight_quantizer: BaseWeightQuantizer object for gradient based weight quantizer
-        """
-        super().__init__()
-
-        # loading operation
-        self.op = node.type(**node.framework_attr)
-
-        # loading the weights from the graph node (weights of the trained model)
-        self.op.load_state_dict({k: torch.Tensor(v) for k, v in node.weights.items()}, strict=False)
-        self.float_weight = to_torch_tensor(getattr(self.op, KERNEL)).detach()
-
-        # replace non-gradient needed nn.Parameter with gradient needed torch.tensor
-        delattr(self.op, KERNEL)
-        setattr(self.op, KERNEL, self.float_weight)
-        setattr(getattr(self.op, KERNEL), 'requires_grad', True)
-
-        # quantizer
-        self.weight_quantizer = weight_quantizer(node.final_weights_quantization_cfg, gptq_config, self.float_weight)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Weight fake quantizer wrapper
-        Args:
-            x: input to layer.
-        Returns:
-            Output of layer after using operation with fake quantized weights
-        """
-        # Run weight quantizer
-        setattr(self.op, KERNEL, self.weight_quantizer(self.float_weight))
-        # Do computation
-        return self.op(x)
-
-
-def quantizer_wrapper(node: BaseNode, gptq_config: GradientPTQConfig) -> nn.Module:
-    """
-    Construct a Pytorch model that constitutes as a wrapper for a Pytorch layer, built from a given graph node.
-    Args:
-        node: Node to build its Pytorch layer.
-        gptq_config: GradientPTQConfig with parameters about the tuning process.
-    """
-    if node.is_weights_quantization_enabled():
-        quantization_method = node.final_weights_quantization_cfg.weights_quantization_method
-        if quantization_method in [QuantizationMethod.SYMMETRIC, QuantizationMethod.POWER_OF_TWO]:
-            # STE quantizer
-            # ---------------
-            if gptq_config.rounding_type == RoundingType.STE:
-                node_instance = WeightQuantizerWrapper(node, gptq_config, STEWeightQuantizer)
-
-            # Symmetric Gumbel rounding quantizer
-            # ------------------------------------
-            elif gptq_config.rounding_type == RoundingType.GumbelRounding:
-                node_instance = WeightQuantizerWrapper(node, gptq_config, SymmetricGumbelWeightQuantizer)
-
-        elif quantization_method == QuantizationMethod.UNIFORM:
-            # Uniform Gumbel rounding quantizer
-            # ------------------------------------
-            if gptq_config.rounding_type == RoundingType.GumbelRounding:
-                node_instance = WeightQuantizerWrapper(node, gptq_config, UniformGumbelWeightQuantizer)
-            else:
-                Logger.error(f"For quantization method {quantization_method}, GPTQ Rounding type {gptq_config.rounding_type} is not supported")
-        else:
-            Logger.error(f"For quantization method {quantization_method}, GPTQ Rounding type {gptq_config.rounding_type} is not supported")
-    else:
-        # No quantization
-        node_instance = node_builder(node)
-
-    return node_instance

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

@@ -1,103 +0,0 @@
-# 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 torch
-import torch.nn as nn
-from typing import List, Union
-from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig
-from model_compression_toolkit.gptq.pytorch.quantizer.gptq_quantizer import BaseWeightQuantizer
-from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
-from model_compression_toolkit.gptq.pytorch.quantizer.quant_utils import ste_round, ste_clip
-from model_compression_toolkit.gptq.common.gptq_constants import AUXVAR
-from model_compression_toolkit.core.common.quantization.node_quantization_config import NodeWeightsQuantizationConfig
-from model_compression_toolkit.core.common.constants import THRESHOLD
-
-
-class STEWeightQuantizer(BaseWeightQuantizer):
-    """
-    Class that implements a quantizer with trainable parameters to be used for GPTQ training.
-    """
-
-    def __init__(self,
-                 weights_quantization_cfg: NodeWeightsQuantizationConfig,
-                 gptq_config: GradientPTQConfig,
-                 weight: torch.nn.Parameter):
-        """
-        Construct a Pytorch model that utilize a fake weight quantizer of STE (Straight Through Estimator) for symmetric quantizer.
-        Args:
-            weights_quantization_cfg: Configuration of weight quantization
-            gptq_config: GradientPTQConfig object with parameters about the tuning process.
-            weight: weight for auxiliary tensor creation.
-        """
-        super().__init__()
-
-        self.signed = True
-        self.num_bits = weights_quantization_cfg.weights_n_bits
-        self.min_int = -int(self.signed) * (2 ** (self.num_bits - int(self.signed)))
-        self.max_int = (2 ** (self.num_bits - int(self.signed))) - 1
-        self.weight_shape = weight.shape
-        self.threshold_values = weights_quantization_cfg.weights_quantization_params.get(THRESHOLD)
-        self.delta_tensor = self.threshold_values / (2 ** (self.num_bits-int(self.signed)))
-        self.max_delta_change = gptq_config.lsb_change_per_bit_width.get(self.num_bits)
-
-        # Set trainable tensors
-        self.set_trainable_params()
-
-        # Create tensors
-        self.delta_tensor = to_torch_tensor(self.delta_tensor)
-        self.max_tensor_change = self.delta_tensor * self.max_delta_change
-
-    def set_trainable_params(self):
-        """
-        A function to set a list of trainable parameters of the quantizer for GPTQ retraining
-        """
-        self.aux_tensor = nn.Parameter(to_torch_tensor(torch.zeros(self.weight_shape)), requires_grad=True)
-        self.trainable_params.update({AUXVAR: self.aux_tensor})
-
-    def get_aux_variable(self) -> torch.Tensor:
-        """
-        Returns auxiliary trainable variables
-        """
-        return self.trainable_params.get(AUXVAR)
-
-    def get_quantization_variable(self) -> Union[torch.Tensor, List]:
-        """
-        Returns quantization trainable variables
-        """
-        return []
-
-    def get_weight_quantization_params(self) -> dict:
-        """
-        Returns weight quantization dictionary params
-        """
-        return {THRESHOLD: self.threshold_values}
-
-    def forward(self, w: nn.Parameter, training: bool = True) -> nn.Parameter:
-        """
-        Weight fake quantizer
-        Args:
-            w: weights to quantize.
-            training: whether in training mode or not
-        Returns:
-            quantized weights
-        """
-        v0 = ste_clip(self.aux_tensor, min_val=-self.max_tensor_change, max_val=self.max_tensor_change)
-        v1 = v0 / self.delta_tensor
-        w0 = torch.round(w / self.delta_tensor).detach()
-        w1 = w0 + v1
-        w2 = ste_round(w1)
-        w3 = ste_clip(w2, min_val=self.min_int, max_val=self.max_int)
-        w_q = self.delta_tensor * w3
-        return w_q
-