mct-nightly 1.11.0.20240304.post404__py3-none-any.whl → 1.11.0.20240306.post426__py3-none-any.whl

model_compression_toolkit/core/keras/pruning/pruning_keras_implementation.py

@@ -44,9 +44,9 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
         Prunes the entry node of a model in Keras.
 
         Args:
-            node: The entry node to be pruned.
-            output_mask: A numpy array representing the mask to be applied to the output channels.
-            fw_info: Framework-specific information object.
+            node (BaseNode): The entry node to be pruned.
+            output_mask (np.ndarray): A numpy array representing the mask to be applied to the output channels.
+            fw_info (FrameworkInfo): Framework-specific information object.
 
         """
         return _prune_keras_edge_node(node=node,
@@ -63,10 +63,10 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
         Prunes an intermediate node in a Keras model.
 
         Args:
-            node: The intermediate node to be pruned.
-            input_mask: A numpy array representing the mask to be applied to the input channels.
-            output_mask: A numpy array representing the mask to be applied to the output channels.
-            fw_info: Framework-specific information object.
+            node (BaseNode): The intermediate node to be pruned.
+            input_mask (np.ndarray): A numpy array representing the mask to be applied to the input channels.
+            output_mask (np.ndarray): A numpy array representing the mask to be applied to the output channels.
+            fw_info (FrameworkInfo): Framework-specific information object.
 
         """
         _edit_node_input_shape(input_mask, node)
@@ -85,9 +85,9 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
         Prunes the exit node of a model in Keras.
 
         Args:
-            node: The exit node to be pruned.
-            input_mask: A numpy array representing the mask to be applied to the input channels.
-            fw_info: Framework-specific information object.
+            node (BaseNode): The exit node to be pruned.
+            input_mask (np.ndarray): A numpy array representing the mask to be applied to the input channels.
+            fw_info (FrameworkInfo): Framework-specific information object.
 
         """
         return _prune_keras_edge_node(node=node,
@@ -100,10 +100,10 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
         Determines whether a node is an entry node in a Keras model.
 
         Args:
-            node: The node to be checked.
+            node (BaseNode): The node to be checked.
 
         Returns:
-            Boolean indicating if the node is an entry node.
+            bool: Boolean indicating if the node is an entry node.
         """
         return _is_keras_node_pruning_section_edge(node)
 
@@ -115,26 +115,26 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
         Determines whether a node is an exit node in a Keras model.
 
         Args:
-            node: The node to be checked.
-            corresponding_entry_node: The entry node of the pruning section that is checked.
-            fw_info: Framework-specific information object.
+            node (BaseNode): The node to be checked.
+            corresponding_entry_node (BaseNode): The entry node of the pruning section that is checked.
+            fw_info (FrameworkInfo): Framework-specific information object.
 
         Returns:
-            Boolean indicating if the node is an exit node.
+            bool: Boolean indicating if the node is an exit node.
         """
         return _is_keras_node_pruning_section_edge(node) and PruningSection.has_matching_channel_count(node,
                                                                                                        corresponding_entry_node,
                                                                                                        fw_info)
 
-    def is_node_intermediate_pruning_section(self, node) -> bool:
+    def is_node_intermediate_pruning_section(self, node: BaseNode) -> bool:
         """
         Determines whether a node is part of the intermediate section in the pruning process of a Keras model.
 
         Args:
-            node: The node to be checked.
+            node (BaseNode): The node to be checked.
 
         Returns:
-            Boolean indicating if the node is part of the intermediate pruning section.
+            bool: Boolean indicating if the node is part of the intermediate pruning section.
         """
         # Nodes that are not Conv2D, Conv2DTranspose, DepthwiseConv2D, or Dense are considered intermediate.
         return node.type not in [keras.layers.DepthwiseConv2D,

model_compression_toolkit/core/pytorch/constants.py

@@ -62,6 +62,9 @@ DIM = 'dim'
 IN_CHANNELS = 'in_channels'
 OUT_CHANNELS = 'out_channels'
 NUM_FEATURES = 'num_features'
+NUM_PARAMETERS = 'num_parameters'
+IN_FEATURES = 'in_features'
+OUT_FEATURES = 'out_features'
 
 # torch devices
 CUDA = 'cuda'

model_compression_toolkit/core/pytorch/kpi_data_facade.py

@@ -22,7 +22,7 @@ from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi_data import compute_kpi_data
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
-from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig
 from model_compression_toolkit.constants import FOUND_TORCH
 
 if FOUND_TORCH:
@@ -38,7 +38,7 @@ if FOUND_TORCH:
 
     def pytorch_kpi_data(in_model: Module,
                          representative_data_gen: Callable,
-                         core_config: CoreConfig = CoreConfig(), # TODO: Why pytorch is initilized and keras not?
+                         core_config: CoreConfig = CoreConfig(mixed_precision_config=MixedPrecisionQuantizationConfig()),
                          fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO,
                          target_platform_capabilities: TargetPlatformCapabilities = PYTORCH_DEFAULT_TPC) -> KPI:
         """
@@ -75,9 +75,9 @@ if FOUND_TORCH:
 
         """
 
-        if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfigV2):
-            Logger.error("KPI data computation can't be executed without MixedPrecisionQuantizationConfigV2 object."
-                         "Given quant_config is not of type MixedPrecisionQuantizationConfigV2.")
+        if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfig):
+            Logger.error("KPI data computation can't be executed without MixedPrecisionQuantizationConfig object."
+                         "Given quant_config is not of type MixedPrecisionQuantizationConfig.")
 
         fw_impl = PytorchImplementation()
 

model_compression_toolkit/core/pytorch/pruning/__init__.py

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

model_compression_toolkit/core/pytorch/pruning/pruning_pytorch_implementation.py

@@ -0,0 +1,315 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import Tuple, Dict
+
+from model_compression_toolkit.core.common.pruning.pruning_framework_implementation import \
+    PruningFrameworkImplementation
+from model_compression_toolkit.core.common.pruning.pruning_section import PruningSection
+from model_compression_toolkit.core.pytorch.pytorch_implementation import PytorchImplementation
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+from model_compression_toolkit.core.common import BaseNode
+from model_compression_toolkit.core.pytorch.constants import BIAS, GROUPS, OUT_CHANNELS, OUT_FEATURES, NUM_FEATURES, \
+    IN_CHANNELS, IN_FEATURES, NUM_PARAMETERS
+import torch
+
+import numpy as np
+
+from model_compression_toolkit.logger import Logger
+
+
+class PruningPytorchImplementation(PytorchImplementation, PruningFrameworkImplementation):
+    """
+    Implementation of the PruningFramework for the Pytorch framework. This class provides
+    concrete implementations of the abstract methods defined in PruningFrameworkImplementation
+    for the Pytorch framework.
+    """
+
+    def prune_entry_node(self,
+                         node: BaseNode,
+                         output_mask: np.ndarray,
+                         fw_info: FrameworkInfo):
+        """
+        Prunes the entry node of a model in Pytorch.
+
+        Args:
+            node (BaseNode): The entry node to be pruned.
+            output_mask (np.ndarray): A numpy array representing the mask to be applied to the output channels.
+            fw_info (FrameworkInfo): Framework-specific information object.
+
+        """
+        return _prune_pytorch_edge_node(node=node,
+                                        mask=output_mask,
+                                        fw_info=fw_info,
+                                        is_exit_node=False)
+
+    def prune_intermediate_node(self,
+                                node: BaseNode,
+                                input_mask: np.ndarray,
+                                output_mask: np.ndarray,
+                                fw_info: FrameworkInfo):
+        """
+        Prunes an intermediate node in a Pytorch model.
+
+        Args:
+            node (BaseNode): The intermediate node to be pruned.
+            input_mask (np.ndarray): A numpy array representing the mask to be applied to the input channels.
+            output_mask (np.ndarray): A numpy array representing the mask to be applied to the output channels.
+            fw_info (FrameworkInfo): Framework-specific information object.
+
+        """
+        # TODO (reuvenp/liord): Address handling of node parameters that can be either a single value across all channels or distinct per channel, e.g., PReLU. Consider developing a structured approach.
+        pruning_en = True
+        _edit_node_input_shape(node, input_mask, fw_info)
+        pruned_parameters = {}
+        mask_bool = output_mask.astype(bool)
+        node.weights = pruned_parameters
+        if node.type == torch.nn.BatchNorm2d:
+            node.framework_attr[NUM_FEATURES] = int(np.sum(input_mask))
+        elif node.type == torch.nn.PReLU:
+            if node.framework_attr[NUM_PARAMETERS] > 1:
+                node.framework_attr[NUM_PARAMETERS] = int(np.sum(input_mask))
+            else:
+                pruning_en = False
+
+        if pruning_en:
+            for k, v in node.weights.items():
+                # Apply the mask to the weights.
+                pruned_parameters[k] = v.compress(mask_bool, axis=-1)
+
+    def prune_exit_node(self,
+                        node: BaseNode,
+                        input_mask: np.ndarray,
+                        fw_info: FrameworkInfo):
+        """
+        Prunes the exit node of a model in Pytorch.
+
+        Args:
+            node (BaseNode): The exit node to be pruned.
+            input_mask (np.ndarray): A numpy array representing the mask to be applied to the input channels.
+            fw_info (FrameworkInfo): Framework-specific information object.
+
+        """
+        return _prune_pytorch_edge_node(node=node,
+                                        mask=input_mask,
+                                        fw_info=fw_info,
+                                        is_exit_node=True)
+
+    def is_node_entry_node(self, node: BaseNode) -> bool:
+        """
+        Determines whether a node is an entry node in a Pytorch model.
+
+        Args:
+            node (BaseNode): The node to be checked.
+
+        Returns:
+            bool: Boolean indicating if the node is an entry node.
+        """
+        return _is_pytorch_node_pruning_section_edge(node)
+
+    def is_node_exit_node(self,
+                          node: BaseNode,
+                          corresponding_entry_node: BaseNode,
+                          fw_info: FrameworkInfo) -> bool:
+        """
+        Determines whether a node is an exit node in a Pytorch model.
+
+        Args:
+            node (BaseNode): The node to be checked.
+            corresponding_entry_node (BaseNode): The entry node of the pruning section that is checked.
+            fw_info (FrameworkInfo) Framework-specific information object.
+
+        Returns:
+            bool: Boolean indicating if the node is an exit node.
+        """
+        return _is_pytorch_node_pruning_section_edge(node) and PruningSection.has_matching_channel_count(node,
+                                                                                                         corresponding_entry_node,
+                                                                                                         fw_info)
+
+    def is_node_intermediate_pruning_section(self, node: BaseNode) -> bool:
+        """
+        Determines whether a node is part of the intermediate section in the pruning process of a Pytorch model.
+
+        Args:
+            node (BaseNode): The node to be checked.
+
+        Returns:
+            bool: Boolean indicating if the node is part of the intermediate pruning section.
+        """
+        # Nodes that are not Conv2d, ConvTranspose2d, or Linear are considered intermediate.
+        # For PReLU prune attributes only if there is a parameter per channel
+        return node.type not in [torch.nn.Conv2d,
+                                 torch.nn.ConvTranspose2d,
+                                 torch.nn.Linear]
+
+    def attrs_oi_channels_info_for_pruning(self,
+                                           node: BaseNode,
+                                           fw_info: FrameworkInfo) -> Dict[str, Tuple[int, int]]:
+        """
+        Retrieves the attributes of a given node along with the output/input (OI) channel axis
+        for each attribute used to prune these attributes.
+
+        Not all attributes of a node are directly associated with both input and output channels.
+        For example, bias vectors in convolutional layers are solely related to the number of output
+        channels and do not have a corresponding input channel dimension.
+        In cases like that, None is returned in the tuple of axis for such attributes.
+
+        For kernel operations (like convolutions), the function identifies the output and input
+        channel axis based on framework-specific information.
+        For non-kernel operations, it defaults to setting the last axis as the output
+        channel axis, assuming no specific input channel axis.
+
+        Args:
+            node (BaseNode): The node from the computational graph.
+            fw_info (FrameworkInfo): Contains framework-specific information and utilities.
+
+        Returns:
+            Dict[str, Tuple[int, int]]: A dictionary where each key is an attribute name (like 'weight' or 'bias')
+            and each value is a tuple representing the output and input channel axis indices respectively.
+        """
+
+        attributes_with_axis = {}
+        if fw_info.is_kernel_op(node.type):
+            kernel_attributes = fw_info.get_kernel_op_attributes(node.type)
+            if kernel_attributes is None or len(kernel_attributes) == 0:
+                Logger.error(f"Expected to find attributes but found {kernel_attributes}")
+
+            for attr in kernel_attributes:
+                attributes_with_axis[attr] = fw_info.kernel_channels_mapping.get(node.type)
+
+            # Bias is a vector at the length of the number of output channels.
+            # For this reason, input channel axis is irrelevant to the bias attribute.
+            attributes_with_axis[BIAS] = (0, None)
+        else:
+            # We have several assumptions here:
+            # 1. For intermediate nodes, we prune all nodes' weights.
+            # 2. The output channel axis is the last axis of this attribute.
+            # 3. The input channel axis is irrelevant since these attributes are pruned only by
+            #    their output channels.
+            for attr in list(node.weights.keys()):
+                # If the number of float parameters is 1 or less - is the case where
+                # we have one parameter for all channels. For this case, we don't
+                # want to prune the parameter.
+                if node.get_num_parameters(fw_info)[1] <= 1:
+                    attributes_with_axis[attr] = (None, None)
+                else:
+                    attributes_with_axis[attr] = (-1, None)
+
+        return attributes_with_axis
+
+
+def _is_pytorch_node_pruning_section_edge(node: BaseNode) -> bool:
+    """
+    Determines if a Pytorch node is an edge of a pruning section.
+
+    In the context of pruning, an 'edge' node is a layer that can potentially be pruned.
+    This function identifies such nodes based on their type and attributes. Specifically,
+    Conv2d and ConvTranspose2d layers with 'groups' attribute set to 1, and Linear layers
+    are considered as edges for pruning sections.
+
+    Args:
+        node (BaseNode): The node to be evaluated.
+
+    Returns:
+        bool: True if the node is an edge of a pruning section, False otherwise.
+    """
+
+    # Check if the node is a Conv2D or Conv2DTranspose layer with groups set to 1.
+    if node.type in [torch.nn.Conv2d, torch.nn.ConvTranspose2d]:
+        return node.framework_attr[GROUPS] == 1
+    return node.type == torch.nn.Linear
+
+
+def _prune_pytorch_edge_node(node: BaseNode,
+                             mask: np.ndarray,
+                             fw_info: FrameworkInfo,
+                             is_exit_node: bool):
+    """
+    Prunes the given Pytorch node by applying the mask to the node's weights (weights and biases).
+    This function can handle both entry and exit nodes by specifying the is_exit_node parameter.
+
+    Args:
+        node (BaseNode): The node to be pruned.
+        mask (np.ndarray): The pruning mask to be applied.
+        fw_info (FrameworkInfo): Framework-specific information object.
+        is_exit_node (bool): A boolean indicating whether the node is an exit node.
+
+    """
+
+    # Retrieve the kernel attribute and the axes to prune.
+    kernel_attr = fw_info.get_kernel_op_attributes(node.type)[0]
+    io_axis = fw_info.kernel_channels_mapping.get(node.type)
+    axis_to_prune = io_axis[int(is_exit_node)]
+    kernel = node.get_weights_by_keys(kernel_attr)
+    # Convert mask to boolean.
+    mask_bool = mask.astype(bool)
+
+    pruned_kernel = kernel.compress(mask_bool, axis=axis_to_prune)
+    node.set_weights_by_keys(name=kernel_attr, tensor=pruned_kernel)
+
+    if not is_exit_node and node.framework_attr[BIAS]:
+        # Prune the bias if applicable and it's an entry node.
+        bias = node.get_weights_by_keys(BIAS)
+        pruned_bias = bias.compress(mask_bool)
+        node.set_weights_by_keys(name=BIAS, tensor=pruned_bias)
+
+    if not is_exit_node:
+        # Update 'out_channels' or 'out_features' attributes for entry nodes
+        # Conv2d,ConvTranspose2d / Linear layers.
+        if node.type in [torch.nn.Conv2d, torch.nn.ConvTranspose2d]:
+            node.framework_attr[OUT_CHANNELS] = int(np.sum(mask))
+        elif node.type == torch.nn.Linear:
+            node.framework_attr[OUT_FEATURES] = int(np.sum(mask))
+        else:
+            Logger.exception(f"{node.type} is currently not supported"
+                             f"as an edge node in a pruning section")
+
+    if is_exit_node:
+        if node.type in [torch.nn.Conv2d, torch.nn.ConvTranspose2d]:
+            node.framework_attr[IN_CHANNELS] = int(np.sum(mask))
+        elif node.type == torch.nn.Linear:
+            node.framework_attr[IN_FEATURES] = int(np.sum(mask))
+        else:
+            Logger.exception(f"{node.type} is currently not supported"
+                             f"as an edge node in a pruning section")
+        # Adjust the input shape for the last node in the section.
+        _edit_node_input_shape(node, mask_bool, fw_info)
+
+
+def _edit_node_input_shape(node: BaseNode,
+                           input_mask: np.ndarray,
+                           fw_info: FrameworkInfo):
+    """
+    Adjusts the input shape of a node based on the given input mask.
+
+    This function modifies the input shape of the given node to reflect the pruning
+    that has taken place. It updates the last dimension of the node's input shape
+    to match the number of channels that remain after pruning.
+
+    Args:
+        node (BaseNode): The node whose input shape needs to be adjusted.
+        input_mask (np.ndarray): A binary array where 1 indicates the channel is kept and 0 means pruned.
+        fw_info (FrameworkInfo): Framework-specific information object.
+    """
+    # Start with the current input shape of the node.
+    new_input_shape = list(node.input_shape)
+
+    # Adjust the last dimension of the shape to match the number of unpruned (retained) channels.
+    # This is done by summing the mask, as each '1' in the mask represents a retained channel.
+    channel_axis = fw_info.out_channel_axis_mapping.get(node.type)
+    new_input_shape[0][channel_axis] = int(np.sum(input_mask))
+
+    # Update the node's input shape with the new dimensions.
+    node.input_shape = tuple(new_input_shape)

model_compression_toolkit/core/pytorch/pytorch_implementation.py

@@ -26,7 +26,7 @@ from torch.nn import Module, Sigmoid, Softmax
 
 import model_compression_toolkit.core.pytorch.constants as pytorch_constants
 from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS
-from model_compression_toolkit.core import QuantizationConfig, FrameworkInfo, CoreConfig, MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core import QuantizationConfig, FrameworkInfo, CoreConfig, MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
@@ -332,7 +332,7 @@ class PytorchImplementation(FrameworkImplementation):
 
     def get_sensitivity_evaluator(self,
                                   graph: Graph,
-                                  quant_config: MixedPrecisionQuantizationConfigV2,
+                                  quant_config: MixedPrecisionQuantizationConfig,
                                   representative_data_gen: Callable,
                                   fw_info: FrameworkInfo,
                                   disable_activation_for_metric: bool = False,

model_compression_toolkit/gptq/__init__.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig, RoundingType, GradientPTQConfigV2, GPTQHessianScoresConfig
+from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig, RoundingType, GradientPTQConfig, GPTQHessianScoresConfig
 from model_compression_toolkit.gptq.keras.quantization_facade import keras_gradient_post_training_quantization
 from model_compression_toolkit.gptq.keras.quantization_facade import get_keras_gptq_config
 from model_compression_toolkit.gptq.pytorch.quantization_facade import pytorch_gradient_post_training_quantization

model_compression_toolkit/gptq/common/gptq_config.py

@@ -61,8 +61,8 @@ class GradientPTQConfig:
     """
     Configuration to use for quantization with GradientPTQ.
     """
-
-    def __init__(self, n_iter: int,
+    def __init__(self,
+                 n_epochs: int,
                  optimizer: Any,
                  optimizer_rest: Any = None,
                  loss: Callable = None,
@@ -79,7 +79,7 @@ class GradientPTQConfig:
         Initialize a GradientPTQConfig.
 
         Args:
-            n_iter (int): Number of iterations to train.
+            n_epochs (int): Number of representative dataset epochs to train.
             optimizer (Any): Optimizer to use.
             optimizer_rest (Any): Optimizer to use for bias and quantizer parameters.
             loss (Callable): The loss to use. should accept 6 lists of tensors. 1st list of quantized tensors, the 2nd list is the float tensors,
@@ -96,7 +96,8 @@ class GradientPTQConfig:
             gptq_quantizer_params_override (dict): A dictionary of parameters to override in GPTQ quantizer instantiation. Defaults to None (no parameters).
 
         """
-        self.n_iter = n_iter
+
+        self.n_epochs = n_epochs
         self.optimizer = optimizer
         self.optimizer_rest = optimizer_rest
         self.loss = loss
@@ -114,71 +115,3 @@ class GradientPTQConfig:
             else gptq_quantizer_params_override
 
 
-class GradientPTQConfigV2(GradientPTQConfig):
-    """
-    Configuration to use for quantization with GradientPTQV2.
-    """
-    def __init__(self, n_epochs: int,
-                 optimizer: Any,
-                 optimizer_rest: Any = None,
-                 loss: Callable = None,
-                 log_function: Callable = None,
-                 train_bias: bool = True,
-                 rounding_type: RoundingType = RoundingType.SoftQuantizer,
-                 use_hessian_based_weights: bool = True,
-                 optimizer_quantization_parameter: Any = None,
-                 optimizer_bias: Any = None,
-                 regularization_factor: float = REG_DEFAULT,
-                 hessian_weights_config: GPTQHessianScoresConfig = GPTQHessianScoresConfig(),
-                 gptq_quantizer_params_override: Dict[str, Any] = None):
-        """
-        Initialize a GradientPTQConfigV2.
-
-        Args:
-            n_epochs (int): Number of representative dataset epochs to train.
-            optimizer (Any): Optimizer to use.
-            optimizer_rest (Any): Optimizer to use for bias and quantizer parameters.
-            loss (Callable): The loss to use. should accept 6 lists of tensors. 1st list of quantized tensors, the 2nd list is the float tensors,
-             the 3rd is a list of quantized weights, the 4th is a list of float weights, the 5th and 6th lists are the mean and std of the tensors
-             accordingly. see example in multiple_tensors_mse_loss
-            log_function (Callable): Function to log information about the GPTQ process.
-            train_bias (bool): Whether to update the bias during the training or not.
-            rounding_type (RoundingType): An enum that defines the rounding type.
-            use_hessian_based_weights (bool): Whether to use Hessian-based weights for weighted average loss.
-            optimizer_quantization_parameter (Any): Optimizer to override the rest optimizer  for quantizer parameters.
-            optimizer_bias (Any): Optimizer to override the rest optimizerfor bias.
-            regularization_factor (float): A floating point number that defines the regularization factor.
-            hessian_weights_config (GPTQHessianScoresConfig): A configuration that include all necessary arguments to run a computation of Hessian scores for the GPTQ loss.
-            gptq_quantizer_params_override (dict): A dictionary of parameters to override in GPTQ quantizer instantiation. Defaults to None (no parameters).
-
-        """
-
-        super().__init__(n_iter=None,
-                         optimizer=optimizer,
-                         optimizer_rest=optimizer_rest,
-                         loss=loss,
-                         log_function=log_function,
-                         train_bias=train_bias,
-                         rounding_type=rounding_type,
-                         use_hessian_based_weights=use_hessian_based_weights,
-                         optimizer_quantization_parameter=optimizer_quantization_parameter,
-                         optimizer_bias=optimizer_bias,
-                         regularization_factor=regularization_factor,
-                         hessian_weights_config=hessian_weights_config,
-                         gptq_quantizer_params_override=gptq_quantizer_params_override)
-        self.n_epochs = n_epochs
-
-    @classmethod
-    def from_v1(cls, n_ptq_iter: int, config_v1: GradientPTQConfig):
-        """
-        Initialize a GradientPTQConfigV2 from GradientPTQConfig instance.
-
-        Args:
-            n_ptq_iter (int): Number of PTQ calibration iters (length of representative dataset).
-            config_v1 (GradientPTQConfig): A GPTQ config to convert to V2.
-
-        """
-        n_epochs = int(round(config_v1.n_iter) / n_ptq_iter)
-        v1_params = config_v1.__dict__
-        v1_params = {k: v for k, v in v1_params.items() if k != 'n_iter'}
-        return cls(n_epochs, **v1_params)

model_compression_toolkit/gptq/keras/gptq_training.py

@@ -37,7 +37,7 @@ else:
 from model_compression_toolkit.trainable_infrastructure import KerasTrainableQuantizationWrapper
 from model_compression_toolkit.core import common
 from model_compression_toolkit.gptq.common.gptq_training import GPTQTrainer
-from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfigV2
+from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.gptq.keras.graph_info import get_weights_for_loss, get_gptq_trainable_parameters
 from model_compression_toolkit.gptq.keras.quantizer.regularization_factory import get_regularization
@@ -56,7 +56,7 @@ class KerasGPTQTrainer(GPTQTrainer):
     def __init__(self,
                  graph_float: Graph,
                  graph_quant: Graph,
-                 gptq_config: GradientPTQConfigV2,
+                 gptq_config: GradientPTQConfig,
                  fw_impl: FrameworkImplementation,
                  fw_info: FrameworkInfo,
                  representative_data_gen: Callable,