mct-nightly 2.4.0.20250925.543__py3-none-any.whl → 2.4.2.20250926.532__py3-none-any.whl

model_compression_toolkit/core/keras/graph_substitutions/substitutions/shift_negative_activation.py

@@ -34,7 +34,6 @@ from model_compression_toolkit.core.common.graph.graph_matchers import NodeOpera
     NodeFrameworkAttrMatcher
 from model_compression_toolkit.core.common.substitutions.shift_negative_activation import \
     apply_shift_negative_correction
-from model_compression_toolkit.core.keras.quantization.activation_quantization_fn_factory import get_activation_quantization_fn_factory
 from model_compression_toolkit.core.keras.constants import KERNEL_SIZE, STRIDES, ACTIVATION, SWISH, \
     SELU, GELU, FUNCTION, ADD, PAD
 from model_compression_toolkit.core.keras.constants import NEGATIVE_SLOPE, PADDING, PAD_SAME, PAD_VALID, BIAS, USE_BIAS
@@ -228,13 +227,15 @@ def is_padding_node_and_node_has_padding(pad_node_to_consider: BaseNode,
 
 
 def keras_apply_shift_negative_correction(graph: Graph,
-                                          core_config: CoreConfig) -> Graph:
+                                          core_config: CoreConfig,
+                                          fw_info: FrameworkInfo) -> Graph:
     """
     Apply shift negative correction (SNC) on a graph built from a Keras model.
 
     Args:
         graph: Graph to apply SNC on.
         core_config: Quantization configuration.
+        fw_info: FrameworkInfo object with information about the specific framework's module.
 
     Returns:
         Graph after SNC.
@@ -243,6 +244,7 @@ def keras_apply_shift_negative_correction(graph: Graph,
 
     return apply_shift_negative_correction(graph,
                                            core_config,
+                                           fw_info,
                                            snc_node,
                                            linear_node,
                                            bypass_node,
@@ -253,6 +255,5 @@ def keras_apply_shift_negative_correction(graph: Graph,
                                            is_padding_node_and_node_has_padding,
                                            PADDING,
                                            BIAS,
-                                           USE_BIAS,
-                                           get_activation_quantization_fn_factory
+                                           USE_BIAS
                                            )

model_compression_toolkit/core/keras/hessian/weights_hessian_scores_calculator_keras.py

@@ -22,6 +22,7 @@ from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS, MIN_HESS
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianScoresGranularity
 from model_compression_toolkit.core.keras.back2framework.float_model_builder import FloatKerasModelBuilder
+from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.core.keras.hessian.hessian_scores_calculator_keras import HessianScoresCalculatorKeras
 from model_compression_toolkit.logger import Logger
 
@@ -94,11 +95,20 @@ class WeightsHessianScoresCalculatorKeras(HessianScoresCalculatorKeras):
                 for i, ipt_node in enumerate(self.hessian_request.target_nodes):  # Per Interest point weights tensor
 
                     # Check if the target node's layer type is supported.
-                    if not ipt_node.kernel_attr:
+                    if not DEFAULT_KERAS_INFO.is_kernel_op(ipt_node.type):
                         Logger.critical(f"Hessian information with respect to weights is not supported for "
                                         f"{ipt_node.type} layers.")  # pragma: no cover
 
-                    weight_tensor = getattr(model.get_layer(ipt_node.name), ipt_node.kernel_attr)
+                    # Get the weight attributes for the target node type
+                    weight_attributes = DEFAULT_KERAS_INFO.get_kernel_op_attributes(ipt_node.type)
+
+                    # Get the weight tensor for the target node
+                    if len(weight_attributes) != 1:  # pragma: no cover
+                        Logger.critical(
+                            f"Hessian-based scoring with respect to weights is currently supported only for nodes with "
+                            f"a single weight attribute. Found {len(weight_attributes)} attributes.")
+
+                    weight_tensor = getattr(model.get_layer(ipt_node.name), weight_attributes[0])
 
                     if j == 0:
                         # On the first iteration we store the weight_tensor shape for later reshaping the results
@@ -106,7 +116,7 @@ class WeightsHessianScoresCalculatorKeras(HessianScoresCalculatorKeras):
                         tensors_original_shape.append(weight_tensor.shape)
 
                     # Get the output channel index (needed for HessianInfoGranularity.PER_OUTPUT_CHANNEL case)
-                    output_channel_axis = ipt_node.channel_axis.output
+                    output_channel_axis, _ = DEFAULT_KERAS_INFO.kernel_channels_mapping.get(ipt_node.type)
 
                     # Get number of scores that should be calculated by the granularity.
                     num_of_scores = self._get_num_scores_by_granularity(weight_tensor,

model_compression_toolkit/core/keras/keras_implementation.py

@@ -65,6 +65,7 @@ from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.model_builder_mode import ModelBuilderMode
 from model_compression_toolkit.core.common.node_prior_info import NodePriorInfo
+from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.core.keras.graph_substitutions.substitutions.activation_decomposition import \
     ActivationDecomposition
 from model_compression_toolkit.core.keras.graph_substitutions.substitutions.matmul_substitution import \
@@ -174,16 +175,18 @@ class KerasImplementation(FrameworkImplementation):
                       graph: Graph,
                       mode: ModelBuilderMode,
                       append2output: List[Any] = None,
+                      fw_info: FrameworkInfo = DEFAULT_KERAS_INFO,
                       return_float_outputs: bool = False) -> Tuple:
         """
         Build a Keras model from a graph.
-        The mode determines how the model should be built. append2output is a list of Nodes
+        The mode determines how the model should be build. append2output is a list of Nodes
         to set as the model outputs.
 
         Args:
             graph: Graph to build the model from it.
             mode: Mode for how to build the model.
             append2output: List of Nodes to set as the model's outputs.
+            fw_info: FrameworkInfo object with information about the specific framework's model
             return_float_outputs (bool): whether to return outputs before or after quantization nodes (default)
         Returns:
             A tuple with the model and additional relevant supporting objects.
@@ -192,6 +195,7 @@ class KerasImplementation(FrameworkImplementation):
         keras_model_builder = get_keras_model_builder(mode)
         return keras_model_builder(graph=graph,
                                    append2output=append2output,
+                                   fw_info=fw_info,
                                    return_float_outputs=return_float_outputs).build_model()
 
     def run_model_inference(self,
@@ -223,57 +227,65 @@ class KerasImplementation(FrameworkImplementation):
 
     def shift_negative_correction(self,
                                   graph: Graph,
-                                  core_config: CoreConfig) -> Graph:
+                                  core_config: CoreConfig,
+                                  fw_info: FrameworkInfo) -> Graph:
         """
         Apply shift negative correction (SNC) on a graph.
 
         Args:
             graph: Graph to apply SNC on.
             core_config: Quantization configuration.
+            fw_info: FrameworkInfo object with information about the specific framework's model.
 
         Returns:
             Graph after SNC.
         """
         return keras_apply_shift_negative_correction(graph,
-                                                     core_config)
+                                                     core_config,
+                                                     fw_info)
 
     def compute_activation_bias_correction(self,
                                            graph: Graph,
-                                           quant_config: QuantizationConfig):
+                                           quant_config: QuantizationConfig,
+                                           fw_info: FrameworkInfo):
         """
         Compute activation bias correction on a graph.
 
         Args:
             graph: Graph to apply activation bias correction on.
             quant_config: QuantizationConfig of how the model should be quantized.
+            fw_info: FrameworkInfo object with information about the specific framework's model.
 
         Returns:
             Graph after activation bias correction computing.
         """
         return keras_compute_activation_bias_correction_of_graph(graph=graph,
                                                                  quant_config=quant_config,
+                                                                 fw_info=fw_info,
                                                                  fw_impl=self)
 
     def get_substitutions_channel_equalization(self,
-                                               quant_config: QuantizationConfig) -> List[common.BaseSubstitution]:
+                                               quant_config: QuantizationConfig,
+                                               fw_info: FrameworkInfo) -> List[common.BaseSubstitution]:
         """
         Return a list of the framework substitutions used for channel equalization.
 
         Args:
             quant_config: QuantizationConfig to determine which substitutions to return.
+            fw_info: FrameworkInfo object with information about the specific framework's model.
 
         Returns:
             A list of the framework substitutions used after we collect statistics.
         """
         substitutions_list = []
         if quant_config.activation_channel_equalization:
-            substitutions_list.extend([ScaleEqualization(quant_config),
-                                       ScaleEqualizationWithPad(quant_config),
-                                       ScaleEqualizationMidActivation(quant_config),
-                                       ScaleEqualizationMidActivationWithPad(quant_config)])
+            substitutions_list.extend([ScaleEqualization(quant_config, fw_info),
+                                       ScaleEqualizationWithPad(quant_config, fw_info),
+                                       ScaleEqualizationMidActivation(quant_config, fw_info),
+                                       ScaleEqualizationMidActivationWithPad(quant_config, fw_info)])
         return substitutions_list
 
-    def get_substitutions_prepare_graph(self) -> List[common.BaseSubstitution]:
+    def get_substitutions_prepare_graph(self, fw_info: FrameworkInfo = None) -> List[common.BaseSubstitution]:
         """
 
         Returns: A list of the framework substitutions used to prepare the graph.
@@ -357,8 +369,8 @@ class KerasImplementation(FrameworkImplementation):
         if quant_config.softmax_shift:
             substitutions_list.append(keras_softmax_shift())
         if quant_config.input_scaling:
-            substitutions_list.append(InputScaling(quant_config))
-            substitutions_list.append(InputScalingWithPad(quant_config))
+            substitutions_list.append(InputScaling())
+            substitutions_list.append(InputScalingWithPad())
         if quant_config.concat_threshold_update:
             substitutions_list.append(ConcatThresholdUpdate())
         return substitutions_list
@@ -390,19 +402,22 @@ class KerasImplementation(FrameworkImplementation):
 
     def get_node_prior_info(self,
                             node: BaseNode,
+                            fw_info: FrameworkInfo,
                             graph: Graph) -> NodePriorInfo:
         """
         Get a NodePriorInfo object for a node that represents a Keras layer.
 
         Args:
             node: Node to get its prior info.
+            fw_info: Framework specific information needed to create the prior info of the node.
             graph: Graph to check the next node type.
 
         Returns:
             NodePriorInfo with information about the node.
         """
 
-        return create_node_prior_info(node=node, graph=graph)
+        return create_node_prior_info(node=node,
+                                      fw_info=fw_info, graph=graph)
 
     def count_node_for_mixed_precision_interest_points(self, node: BaseNode) -> bool:
         """
@@ -515,19 +530,23 @@ class KerasImplementation(FrameworkImplementation):
         return True
 
     def get_node_mac_operations(self,
-                                node: BaseNode) -> float:
+                                node: BaseNode,
+                                fw_info: FrameworkInfo) -> float:
         """
         Gets the MAC operation count for a given operation.
 
         Args:
             node: A graph node that wraps the operation for which the MAC count is computed.
+            fw_info: FrameworkInfo object with information about the Keras model.
 
         Returns: The MAC count og the operation
         """
-        if node.kernel_attr is None:
+        kernels = fw_info.get_kernel_op_attributes(node.type)
+        if not kernels or kernels[0] is None:
             return 0
 
-        kernel_shape = node.get_weights_by_keys(node.kernel_attr).shape
+        assert len(kernels) == 1
+        kernel_shape = node.get_weights_by_keys(kernels[0]).shape
 
         if node.is_match_type(Conv2D) or node.is_match_type(Conv2DTranspose) or node.is_match_type(DepthwiseConv2D):
             h, w = node.get_output_shapes_list()[0][-3:-1]
@@ -535,7 +554,8 @@ class KerasImplementation(FrameworkImplementation):
 
         if node.is_match_type(Dense):
             # IN * OUT * (all previous dims[:-1])
-            return node.get_total_output_params() * kernel_shape[node.channel_axis.input]
+            _, input_channel_axis = fw_info.kernel_channels_mapping.get(node.type)
+            return node.get_total_output_params() * kernel_shape[input_channel_axis]
 
         return 0
 

model_compression_toolkit/core/keras/keras_model_validation.py

@@ -0,0 +1,38 @@
+from tensorflow.keras.models import Model
+
+from model_compression_toolkit.core import FrameworkInfo
+from model_compression_toolkit.core.common.framework_info import ChannelAxis
+from model_compression_toolkit.core.common.model_validation import ModelValidation
+from model_compression_toolkit.core.keras.constants import CHANNELS_FORMAT, CHANNELS_FORMAT_LAST, CHANNELS_FORMAT_FIRST
+
+
+class KerasModelValidation(ModelValidation):
+    """
+    Class to define validation methods in order to validate the received Keras model to quantize.
+    """
+
+    def __init__(self, model: Model, fw_info: FrameworkInfo):
+        """
+        Initialize a KerasModelValidation object.
+
+        Args:
+            model: Keras model to check its validity.
+            fw_info: Information about the framework of the model (Keras).
+        """
+
+        super(KerasModelValidation, self).__init__(model=model,
+                                                   fw_info=fw_info)
+
+    def validate_output_channel_consistency(self):
+        """
+
+        Validate that output channels index in all layers of the model are the same.
+        If the model has layers with different output channels index, an exception is thrown.
+
+        """
+        for layer in self.model.layers:
+            data_format = layer.get_config().get(CHANNELS_FORMAT)
+            if data_format is not None:
+                assert (data_format == CHANNELS_FORMAT_LAST and self.fw_info.out_channel_axis_mapping.get(layer) == ChannelAxis.NHWC.value
+                        or data_format == CHANNELS_FORMAT_FIRST and self.fw_info.out_channel_axis_mapping.get(layer) == ChannelAxis.NCHW.value), \
+                    f'Model can not have layers with different data formats.'

model_compression_toolkit/core/keras/keras_node_prior_info.py

@@ -17,19 +17,22 @@ from model_compression_toolkit.core.common.graph.base_graph import Graph
 
 
 def create_node_prior_info(node: BaseNode,
+                           fw_info: FrameworkInfo,
                            graph: Graph):
     """
     Create a NodePriorInfo object for a given node.
 
     Args:
         node: Node to create its prior info.
+        fw_info: Information about a specific framework the node was generated from.
         graph: Graph to check the next node type.
 
     Returns:
         NodePriorInfo object with info about the node.
     """
 
-    min_output, max_output = _get_min_max_outputs(node=node)
+    min_output, max_output = _get_min_max_outputs(node=node,
+                                                  fw_info=fw_info)
 
     mean_output, std_output = _get_mean_std_outputs(node=node,
                                                     graph=graph)
@@ -39,12 +42,14 @@ def create_node_prior_info(node: BaseNode,
                          std_output=std_output)
 
 
-def _get_min_max_outputs(node: BaseNode) -> Tuple[Any, Any]:
+def _get_min_max_outputs(node: BaseNode,
+                         fw_info: FrameworkInfo) -> Tuple[Any, Any]:
     """
     Return the min/max output values of a node if known.
     If one of them (or both of them) is unknown - return None instead of a value.
     Args:
         node: Node to create its prior info.
+        fw_info: Information about a specific framework the node was generated from.
 
     Returns:
         Min/max output values if known.
@@ -53,8 +58,12 @@ def _get_min_max_outputs(node: BaseNode) -> Tuple[Any, Any]:
 
     if node.is_match_type(ReLU):
         min_output = node.framework_attr[THRESHOLD] if node.framework_attr[NEGATIVE_SLOPE] == 0 else None
-    else:
-        min_output, max_output = node.minmax
+
+    elif fw_info.layers_has_min_max(node.type):
+        min_output, max_output = fw_info.layer_min_max_mapping[node.type]
+
+    elif node.is_match_type(Activation) and fw_info.activation_has_min_max(node.framework_attr[ACTIVATION]):
+        min_output, max_output = fw_info.activation_min_max_mapping[node.framework_attr[ACTIVATION]]
 
     return min_output, max_output
 

model_compression_toolkit/core/keras/mixed_precision/configurable_activation_quantizer.py

@@ -23,7 +23,6 @@ from model_compression_toolkit.core.common.mixed_precision.configurable_quantize
     verify_candidates_descending_order, init_activation_quantizers
 from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import \
     CandidateNodeQuantizationConfig
-from model_compression_toolkit.core.keras.quantization.activation_quantization_fn_factory import get_activation_quantization_fn_factory
 from model_compression_toolkit.logger import Logger
 
 import tensorflow as tf
@@ -68,7 +67,7 @@ class ConfigurableActivationQuantizer(BaseKerasInferableQuantizer):
             if qc.activation_quantization_cfg.quant_mode != node_q_cfg[0].activation_quantization_cfg.quant_mode:
                 Logger.critical("Unsupported configuration: Mixing candidates with differing activation quantization states (enabled/disabled).")  # pragma: no cover
 
-        self.activation_quantizers = init_activation_quantizers(self.node_q_cfg, get_activation_quantization_fn_factory)
+        self.activation_quantizers = init_activation_quantizers(self.node_q_cfg)
         self.active_quantization_config_index = max_candidate_idx  # initialize with first config as default
 
     def set_active_activation_quantizer(self, index: Optional[int]):

model_compression_toolkit/core/keras/pruning/pruning_keras_implementation.py

@@ -19,6 +19,7 @@ from model_compression_toolkit.core.common.pruning.pruning_framework_implementat
     PruningFrameworkImplementation
 from model_compression_toolkit.core.common.pruning.pruning_section import PruningSection
 from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.keras.constants import BIAS, GROUPS, FILTERS, UNITS, USE_BIAS
 import keras
@@ -28,10 +29,6 @@ import numpy as np
 from model_compression_toolkit.logger import Logger
 
 
-# default output channel axis to use when it's not defined in node's fw_info.
-_default_output_channel_axis = -1
-
-
 class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementation):
     """
     Implementation of the PruningFramework for the Keras framework. This class provides
@@ -41,23 +38,27 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
 
     def prune_entry_node(self,
                          node: BaseNode,
-                         output_mask: np.ndarray):
+                         output_mask: np.ndarray,
+                         fw_info: FrameworkInfo):
         """
         Prunes the entry node of a model in Keras.
 
         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_keras_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):
+                                output_mask: np.ndarray,
+                                fw_info: FrameworkInfo):
         """
         Prunes an intermediate node in a Keras model.
 
@@ -65,6 +66,7 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
             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)
@@ -77,17 +79,20 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
 
     def prune_exit_node(self,
                         node: BaseNode,
-                        input_mask: np.ndarray):
+                        input_mask: np.ndarray,
+                        fw_info: FrameworkInfo):
         """
         Prunes the exit node of a model in Keras.
 
         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_keras_edge_node(node=node,
                                       mask=input_mask,
+                                      fw_info=fw_info,
                                       is_exit_node=True)
 
     def is_node_entry_node(self, node: BaseNode) -> bool:
@@ -104,19 +109,22 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
 
     def is_node_exit_node(self,
                           node: BaseNode,
-                          corresponding_entry_node: BaseNode) -> bool:
+                          corresponding_entry_node: BaseNode,
+                          fw_info: FrameworkInfo) -> bool:
         """
         Determines whether a node is an exit node in a Keras 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_keras_node_pruning_section_edge(node) and PruningSection.has_matching_channel_count(node,
-                                                                                                       corresponding_entry_node)
+                                                                                                       corresponding_entry_node,
+                                                                                                       fw_info)
 
     def is_node_intermediate_pruning_section(self, node: BaseNode) -> bool:
         """
@@ -135,7 +143,8 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
                                  keras.layers.Dense]
 
     def attrs_oi_channels_info_for_pruning(self,
-                                           node: BaseNode) -> Dict[str, Tuple[int, int]]:
+                                           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.
@@ -152,6 +161,7 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
 
         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 'kernel' or 'bias')
@@ -159,8 +169,13 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
         """
 
         attributes_with_axis = {}
-        if node.kernel_attr:
-            attributes_with_axis[node.kernel_attr] = (node.channel_axis.output, node.channel_axis.input)
+        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.critical(f"Expected kernel attributes for operation for node type {node.type}, found None or empty.")
+
+            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.
@@ -176,10 +191,6 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
 
         return attributes_with_axis
 
-    @property
-    def default_output_channel_axis(self):
-        return _default_output_channel_axis
-
 
 def _is_keras_node_pruning_section_edge(node: BaseNode) -> bool:
     """
@@ -205,6 +216,7 @@ def _is_keras_node_pruning_section_edge(node: BaseNode) -> bool:
 
 def _prune_keras_edge_node(node: BaseNode,
                            mask: np.ndarray,
+                           fw_info: FrameworkInfo,
                            is_exit_node: bool):
     """
     Prunes the given Keras node by applying the mask to the node's weights (kernels and biases).
@@ -213,18 +225,21 @@ def _prune_keras_edge_node(node: BaseNode,
     Args:
         node: The node to be pruned.
         mask: The pruning mask to be applied.
+        fw_info: Framework-specific information object.
         is_exit_node: A boolean indicating whether the node is an exit node.
 
     """
 
     # Retrieve the kernel attribute and the axes to prune.
-    axis_to_prune = node.channel_axis.input if is_exit_node else node.channel_axis.output
-    kernel = node.get_weights_by_keys(node.kernel_attr)
+    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=node.kernel_attr, tensor=pruned_kernel)
+    node.set_weights_by_keys(name=kernel_attr, tensor=pruned_kernel)
 
     if not is_exit_node and node.framework_attr[USE_BIAS]:
         # Prune the bias if applicable and it's an entry node.

model_compression_toolkit/core/keras/resource_utilization_data_facade.py

@@ -27,15 +27,14 @@ if FOUND_TF:
     from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2keras import \
         AttachTpcToKeras
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
+    from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
     from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
-    from model_compression_toolkit.core.keras.default_framework_info import set_keras_info
     from tensorflow.keras.models import Model
 
     from model_compression_toolkit import get_target_platform_capabilities
 
     KERAS_DEFAULT_TPC = get_target_platform_capabilities(TENSORFLOW, DEFAULT_TP_MODEL)
 
-    @set_keras_info
     def keras_resource_utilization_data(in_model: Model,
                                         representative_data_gen: Callable,
                                         core_config: CoreConfig = CoreConfig(
@@ -94,6 +93,7 @@ if FOUND_TF:
                                                  representative_data_gen,
                                                  core_config,
                                                  target_platform_capabilities,
+                                                 DEFAULT_KERAS_INFO,
                                                  fw_impl)
 
 else:

model_compression_toolkit/core/keras/statistics_correction/keras_compute_activation_bias_correction_of_graph.py

@@ -25,7 +25,7 @@ else:
 
 from model_compression_toolkit.core import QuantizationConfig
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
-from model_compression_toolkit.core.keras.quantization.activation_quantization_fn_factory import get_activation_quantization_fn_factory
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
 from model_compression_toolkit.core.common.statistics_correction.compute_activation_bias_correction_of_graph import \
@@ -43,6 +43,7 @@ def activation_bias_correction_node_matchers():
 
 def keras_compute_activation_bias_correction_of_graph(graph: Graph,
                                                       quant_config: QuantizationConfig,
+                                                      fw_info: FrameworkInfo,
                                                       fw_impl: FrameworkImplementation) -> Graph:
     """
     Compute the activation bias correction term for graph based on a Keras model.
@@ -50,6 +51,7 @@ def keras_compute_activation_bias_correction_of_graph(graph: Graph,
     Args:
         graph: Graph with nodes to compute the activation bias correction.
         quant_config: QuantizationConfig of how the model should be quantized.
+        fw_info: Framework info like lists of nodes their kernel should quantized.
         fw_impl: FrameworkImplementation object with a specific framework methods implementation.
 
     Returns:
@@ -57,9 +59,9 @@ def keras_compute_activation_bias_correction_of_graph(graph: Graph,
     """
     graph = compute_activation_bias_correction_of_graph(graph=graph,
                                                         quant_config=quant_config,
+                                                        fw_info=fw_info,
                                                         fw_impl=fw_impl,
                                                         activation_bias_correction_node_matchers=
                                                         activation_bias_correction_node_matchers,
-                                                        kernel_size=KERNEL_SIZE,
-                                                        get_activation_quantization_fn_factory=get_activation_quantization_fn_factory)
+                                                        kernel_size=KERNEL_SIZE)
     return graph