mct-nightly 2.4.0.20250617.613__py3-none-any.whl → 2.4.0.20250618.606__py3-none-any.whl

model_compression_toolkit/core/graph_prep_runner.py

@@ -37,7 +37,6 @@ from model_compression_toolkit.target_platform_capabilities.targetplatform2frame
 def graph_preparation_runner(in_model: Any,
                              representative_data_gen: Callable,
                              quantization_config: QuantizationConfig,
-                             fw_info: FrameworkInfo,
                              fw_impl: FrameworkImplementation,
                              fqc: FrameworkQuantizationCapabilities,
                              bit_width_config: BitWidthConfig = None,
@@ -56,8 +55,6 @@ def graph_preparation_runner(in_model: Any,
         in_model (Any): Model to quantize.
         representative_data_gen (Callable): Dataset used for calibration.
         quantization_config (QuantizationConfig): QuantizationConfig containing parameters of how the model should be quantized.
-        fw_info (FrameworkInfo): Information needed for quantization about the specific framework (e.g., kernel channels indices,
-            groups of layers by how they should be quantized, etc.).
         fw_impl (FrameworkImplementation): FrameworkImplementation object with a specific framework methods implementation.
         fqc (FrameworkQuantizationCapabilities): FrameworkQuantizationCapabilities object that models the inference target platform and
             the attached framework operator's information.
@@ -73,7 +70,6 @@ def graph_preparation_runner(in_model: Any,
     graph = read_model_to_graph(in_model,
                                 representative_data_gen,
                                 fqc,
-                                fw_info,
                                 fw_impl)
 
     if tb_w is not None:
@@ -83,7 +79,6 @@ def graph_preparation_runner(in_model: Any,
                                             fqc,
                                             quantization_config,
                                             bit_width_config,
-                                            fw_info,
                                             tb_w,
                                             fw_impl,
                                             mixed_precision_enable=mixed_precision_enable,
@@ -96,7 +91,6 @@ def get_finalized_graph(initial_graph: Graph,
                         fqc: FrameworkQuantizationCapabilities,
                         quant_config: QuantizationConfig = DEFAULTCONFIG,
                         bit_width_config: BitWidthConfig = None,
-                        fw_info: FrameworkInfo = None,
                         tb_w: TensorboardWriter = None,
                         fw_impl: FrameworkImplementation = None,
                         mixed_precision_enable: bool = False,
@@ -111,8 +105,6 @@ def get_finalized_graph(initial_graph: Graph,
         quant_config (QuantizationConfig): QuantizationConfig containing parameters of how the model should be
             quantized.
         bit_width_config (BitWidthConfig): Config for bit-width selection. Defaults to None.
-        fw_info (FrameworkInfo): Information needed for quantization about the specific framework (e.g.,
-            kernel channels indices, groups of layers by how they should be quantized, etc.)
         tb_w (TensorboardWriter): TensorboardWriter object to use for logging events such as graphs, histograms, etc.
         fw_impl (FrameworkImplementation): FrameworkImplementation object with a specific framework methods implementation.
         mixed_precision_enable: is mixed precision enabled.
@@ -124,7 +116,7 @@ def get_finalized_graph(initial_graph: Graph,
     ######################################
     # Graph substitution (prepare graph)
     ######################################
-    graph = substitute(initial_graph, fw_impl.get_substitutions_prepare_graph(fw_info))
+    graph = substitute(initial_graph, fw_impl.get_substitutions_prepare_graph())
 
     if tb_w is not None:
         tb_w.add_graph(graph, 'after_graph_preparation')
@@ -134,7 +126,6 @@ def get_finalized_graph(initial_graph: Graph,
     ##########################################
     for node in graph.nodes:
         node.prior_info = fw_impl.get_node_prior_info(node=node,
-                                                      fw_info=fw_info,
                                                       graph=graph)
 
     ##################################################
@@ -170,8 +161,7 @@ def get_finalized_graph(initial_graph: Graph,
     # Channel equalization
     ######################################
     transformed_graph = substitute(transformed_graph,
-                                   fw_impl.get_substitutions_channel_equalization(quant_config,
-                                                                                  fw_info))
+                                   fw_impl.get_substitutions_channel_equalization(quant_config))
 
     if tb_w is not None:
         tb_w.add_graph(transformed_graph, 'after_graph_marking')
@@ -190,7 +180,6 @@ def get_finalized_graph(initial_graph: Graph,
 def read_model_to_graph(in_model: Any,
                         representative_data_gen: Callable,
                         fqc: FrameworkQuantizationCapabilities,
-                        fw_info: FrameworkInfo = None,
                         fw_impl: FrameworkImplementation = None) -> Graph:
 
     """
@@ -201,8 +190,6 @@ def read_model_to_graph(in_model: Any,
         representative_data_gen: Dataset used for calibration.
         fqc: FrameworkQuantizationCapabilities object that models the inference target platform and
                       the attached framework operator's information.
-        fw_info: Information needed for quantization about the specific framework (e.g.,
-                kernel channels indices, groups of layers by how they should be quantized, etc.)
         fw_impl: FrameworkImplementation object with a specific framework methods implementation.
 
     Returns:
@@ -210,6 +197,5 @@ def read_model_to_graph(in_model: Any,
     """
     graph = fw_impl.model_reader(in_model,
                                  representative_data_gen)
-    graph.set_fw_info(fw_info)
     graph.set_fqc(fqc)
     return graph

model_compression_toolkit/core/keras/back2framework/float_model_builder.py

@@ -17,7 +17,6 @@ from typing import List
 from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder
-from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.core import common
 from tensorflow.python.util.object_identity import Reference as TFReference
 
@@ -29,20 +28,17 @@ class FloatKerasModelBuilder(KerasModelBuilder):
     def __init__(self,
                  graph: common.Graph,
                  append2output=None,
-                 fw_info: FrameworkInfo = DEFAULT_KERAS_INFO,
                  return_float_outputs: bool = False):
         """
 
         Args:
             graph: Graph to build the model from.
             append2output: Nodes to append to model's output.
-            fw_info: Information about the specific framework of the model that is built.
             return_float_outputs: Whether the model returns float tensors or not.
         """
 
         super().__init__(graph,
                          append2output,
-                         fw_info,
                          return_float_outputs)
 
     def _quantize_node_activations(self,

model_compression_toolkit/core/keras/back2framework/keras_model_builder.py

@@ -35,8 +35,6 @@ from typing import Any, Dict, List, Tuple, Callable
 from tensorflow.python.util.object_identity import Reference as TFReference
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core import common
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo
-from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.graph.edge import EDGE_SINK_INDEX
 from model_compression_toolkit.core.keras.back2framework.instance_builder import OperationHandler
@@ -57,7 +55,6 @@ class KerasModelBuilder(BaseModelBuilder):
     def __init__(self,
                  graph: common.Graph,
                  append2output=None,
-                 fw_info: FrameworkInfo = DEFAULT_KERAS_INFO,
                  return_float_outputs: bool = False,
                  wrapper: Callable = None,
                  get_activation_quantizer_holder_fn: Callable=None):
@@ -66,7 +63,6 @@ class KerasModelBuilder(BaseModelBuilder):
         Args:
             graph: Graph to build the model from.
             append2output: Nodes to append to model's output.
-            fw_info: Information about the specific framework of the model that is built.
             return_float_outputs: Whether the model returns float tensors or not.
             wrapper: A function wrapper keras Layers.
             get_activation_quantizer_holder_fn: Function to retrieve a quantization holder for a node.
@@ -75,7 +71,6 @@ class KerasModelBuilder(BaseModelBuilder):
 
         super().__init__(graph,
                          append2output,
-                         fw_info,
                          return_float_outputs)
 
         # Build an OperationHandler to handle conversions from graph nodes to Keras operators.

model_compression_toolkit/core/keras/back2framework/mixed_precision_model_builder.py

@@ -36,7 +36,6 @@ from model_compression_toolkit.core.keras.mixed_precision.configurable_weights_q
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
-from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 
 
 class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
@@ -47,14 +46,12 @@ class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
     def __init__(self,
                  graph: common.Graph,
                  append2output=None,
-                 fw_info: FrameworkInfo = DEFAULT_KERAS_INFO,
                  return_float_outputs: bool = False):
         """
 
         Args:
             graph: Graph to build the model from.
             append2output: Nodes to append to model's output.
-            fw_info: Information about the specific framework of the model that is built.
             return_float_outputs: Whether the model returns float tensors or not.
         """
 
@@ -62,7 +59,6 @@ class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
 
         super().__init__(graph,
                          append2output,
-                         fw_info,
                          return_float_outputs,
                          wrapper=self.mixed_precision_wrapper,
                          get_activation_quantizer_holder_fn=self.mixed_precision_activation_holder)
@@ -87,13 +83,12 @@ class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
             ValueError: if kernel attribute is quantized but not configurable.
         """
 
-        kernel_attr = self.fw_info.get_kernel_op_attributes(n.type)[0]
-        if kernel_attr is None or not n.is_weights_quantization_enabled(kernel_attr):
+        if n.kernel_attr is None or not n.is_weights_quantization_enabled(n.kernel_attr):
             return layer
-        if not n.is_configurable_weight(kernel_attr):  # pragma: no cover
+        if not n.is_configurable_weight(n.kernel_attr):  # pragma: no cover
             raise ValueError(f'Weight wrapper is not expected to be created for non-configurable weight of node {n}.')
-        wq = ConfigurableWeightsQuantizer(**self._get_weights_configurable_quantizer_kwargs(n, kernel_attr))
-        return KerasQuantizationWrapper(layer, weights_quantizers={kernel_attr: wq})
+        wq = ConfigurableWeightsQuantizer(**self._get_weights_configurable_quantizer_kwargs(n, n.kernel_attr))
+        return KerasQuantizationWrapper(layer, weights_quantizers={n.kernel_attr: wq})
 
     def _get_weights_configurable_quantizer_kwargs(self, n: BaseNode, attr: str) -> Dict[str, Any]:
         """
@@ -147,13 +142,12 @@ class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
         # activation number of bits (in reversed order).
         # since only kernel attribute is quantized in weights mixed precision,
         # if the node doesn't have a kernel attribute, we only sort by activation_n_bits.
-        n.sort_node_candidates(self.fw_info)
+        n.sort_node_candidates()
 
         max_candidate_idx = n.find_max_candidate_index()
-        kernel_attr = self.fw_info.get_kernel_op_attributes(n.type)[0]
         activation_quantizers = [ConfigurableActivationQuantizer(**{'node_q_cfg': node_q_cfg_candidates,
                                                                     'max_candidate_idx': max_candidate_idx,
-                                                                    'kernel_attr': kernel_attr})] \
+                                                                    'kernel_attr': n.kernel_attr})] \
                                  * num_of_outputs
 
         # Holder by definition uses a single quantizer for the activation quantization
@@ -181,7 +175,7 @@ class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
 
         # creating a mapping between graph nodes and model's layers for mixed precision configurability
         conf_node2layers = {n.name: self._find_layers_in_model_by_node(n, model.layers)
-                            for n in self.graph.get_configurable_sorted_nodes(self.fw_info)}
+                            for n in self.graph.get_configurable_sorted_nodes()}
 
         return model, user_info, conf_node2layers
 
@@ -231,8 +225,7 @@ class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
 
         """
         # Only layers with kernel op are considered weights configurable
-        kernel_attr = self.fw_info.get_kernel_op_attributes(n.type)[0]
-        weights_quant = False if kernel_attr is None else n.is_weights_quantization_enabled(kernel_attr)
+        weights_quant = False if n.kernel_attr is None else n.is_weights_quantization_enabled(n.kernel_attr)
         act_quant = n.is_activation_quantization_enabled()
 
         if weights_quant and not act_quant:

model_compression_toolkit/core/keras/back2framework/quantized_model_builder.py

@@ -18,7 +18,6 @@ from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder
-from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from tensorflow.python.util.object_identity import Reference as TFReference
 
 
@@ -30,20 +29,17 @@ class QuantizedKerasModelBuilder(KerasModelBuilder):
     def __init__(self,
                  graph: common.Graph,
                  append2output=None,
-                 fw_info: FrameworkInfo = DEFAULT_KERAS_INFO,
                  return_float_outputs: bool = False):
         """
 
         Args:
             graph: Graph to build the model from.
             append2output: Nodes to append to model's output.
-            fw_info: Information about the specific framework of the model that is built.
             return_float_outputs: Whether the model returns float tensors or not.
         """
 
         super().__init__(graph,
                          append2output,
-                         fw_info,
                          return_float_outputs)
 
     def _quantize_node_activations(self,

model_compression_toolkit/core/keras/default_framework_info.py

@@ -13,102 +13,153 @@
 # limitations under the License.
 # ==============================================================================
 
-
 import tensorflow as tf
 
+from typing import Tuple, Any, Dict
+from functools import wraps
+
 from model_compression_toolkit.core.keras.quantizer.lut_fake_quant import activation_lut_kmean_quantizer
 from packaging import version
 
 if version.parse(tf.__version__) >= version.parse("2.13"):
-    from keras.src.layers import Conv2D, DepthwiseConv2D, Dense, Conv2DTranspose, Softmax, ELU
+    from keras.src.layers import Conv2D, DepthwiseConv2D, Dense, Conv2DTranspose, Softmax, ELU, Activation
 else:
-    from keras.layers import Conv2D, DepthwiseConv2D, Dense, Conv2DTranspose, Softmax, ELU  # pragma: no cover
-
-from model_compression_toolkit.defaultdict import DefaultDict
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo, DEFAULT_KERNEL_ATTRIBUTES
+    from keras.layers import Conv2D, DepthwiseConv2D, Dense, Conv2DTranspose, Softmax, ELU, Activation  # pragma: no cover
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo, set_fw_info, ChannelAxisMapping
 from mct_quantizers import QuantizationMethod
-from model_compression_toolkit.constants import SOFTMAX_THRESHOLD
+from model_compression_toolkit.constants import SOFTMAX_THRESHOLD, ACTIVATION
 from model_compression_toolkit.core.keras.constants import SOFTMAX, LINEAR, RELU, SWISH, SIGMOID, IDENTITY, TANH, SELU, \
     KERNEL, DEPTHWISE_KERNEL, GELU
 from model_compression_toolkit.core.keras.quantizer.fake_quant_builder import power_of_two_quantization, symmetric_quantization, uniform_quantization
 
-"""
-Map each layer to a list of its' weights attributes that should get quantized.
-If a layer that is not listed here is queried, [None] is returned.
-"""
-KERNEL_ATTRIBUTES = DefaultDict({Conv2D: [KERNEL],
-                                 DepthwiseConv2D: [DEPTHWISE_KERNEL],
-                                 Dense: [KERNEL],
-                                 Conv2DTranspose: [KERNEL]}, DEFAULT_KERNEL_ATTRIBUTES)
-
-
-"""
-Map a layer to its kernel's output and input channels indices.
-Map's values are tuples of (output_channel_index, input_channel_index).
-Default value is returned for layers that are not included.
-"""
-DEFAULT_CHANNEL_AXIS_DICT = DefaultDict({Conv2D: (3, 2),
-                                         DepthwiseConv2D: (2, 2),
-                                         Dense: (1, 0),
-                                         Conv2DTranspose: (2, 3)}, (None, None))
-
-
-"""
-Map a layer to its output channel axis. 
-Where axis=-1 is the last axis
-"""
-DEFAULT_OUT_CHANNEL_AXIS_DICT = DefaultDict({Conv2D: -1,
-                                             DepthwiseConv2D: -1,
-                                             Dense: -1,
-                                             Conv2DTranspose: -1},
-                                            -1)
-
-
-"""
-Map from an activation function to its min/max output values (if known).
-The values are used for tensor min/max values initialization.
-"""
-ACTIVATION2MINMAX = {SOFTMAX: (0, SOFTMAX_THRESHOLD),
-                     SIGMOID: (0, 1),
-                     LINEAR: (None, None),
-                     IDENTITY: (None, None),
-                     TANH: (-1, 1),
-                     SWISH: (-0.279, None),
-                     RELU: (0, None),
-                     SELU: (-1.76, None),
-                     GELU: (-0.17, None),
-                     }
-
-"""
-Map from an Keras layer to its min/max output values (if known).
-The values are used for tensor min/max values initialization.
-"""
-LAYER2MINMAX = {Softmax: (0, SOFTMAX_THRESHOLD),
-                ELU: (-1, None),
-                tf.nn.silu: (-0.279, None),
-                tf.nn.swish: (-0.279, None),
-                tf.nn.sigmoid: (0, 1),
-                tf.nn.tanh: (-1, 1),
-                tf.nn.relu: (0, None),
-                tf.nn.relu6: (0, None),
-                tf.nn.gelu: (-0.17, None),
-                tf.nn.elu: (-1, None),
-                tf.nn.selu: (-1.76, None),
-                tf.nn.softplus: (0, None),
-                tf.nn.softmax: (0, SOFTMAX_THRESHOLD),
-                }
-"""
-Mapping from a QuantizationMethod to an activation quantizer function.
-"""
-ACTIVATION_QUANTIZER_MAPPING = {QuantizationMethod.POWER_OF_TWO: power_of_two_quantization,
-                                QuantizationMethod.SYMMETRIC: symmetric_quantization,
-                                QuantizationMethod.UNIFORM: uniform_quantization,
-                                QuantizationMethod.LUT_POT_QUANTIZER: activation_lut_kmean_quantizer}
-
-
-DEFAULT_KERAS_INFO = FrameworkInfo(ACTIVATION_QUANTIZER_MAPPING,
-                                   DEFAULT_CHANNEL_AXIS_DICT,
-                                   ACTIVATION2MINMAX,
-                                   LAYER2MINMAX,
-                                   KERNEL_ATTRIBUTES,
-                                   DEFAULT_OUT_CHANNEL_AXIS_DICT)
+
+class KerasInfo(FrameworkInfo):
+    """
+    Extra field defined to handle Activation layer functions:
+
+        _activation_min_max_mapping (Dict[str, tuple]): Dictionary from an activation function to its min/max output values.
+
+    """
+
+    """
+    Map each layer to it's weight attribute that should get quantized.
+    If a layer that is not listed here is queried, None is returned.
+    """
+    kernel_ops_attribute_mapping = {Conv2D: KERNEL,
+                                    DepthwiseConv2D: DEPTHWISE_KERNEL,
+                                    Dense: KERNEL,
+                                    Conv2DTranspose: KERNEL}
+
+    """
+    Map a layer to its kernel's output and input channels indices.
+    Map's values are tuples of (output_channel_index, input_channel_index).
+    Default value is returned for layers that are not included.
+    """
+    kernel_channels_mapping = {Conv2D: ChannelAxisMapping(3, 2),
+                               DepthwiseConv2D: ChannelAxisMapping(2, 2),
+                               Dense: ChannelAxisMapping(1, 0),
+                               Conv2DTranspose: ChannelAxisMapping(2, 3)}
+
+    """
+    Map a layer to its output channel axis.
+    Where axis=-1 is the last axis
+    """
+    out_channel_axis_mapping = {Conv2D: -1,
+                                DepthwiseConv2D: -1,
+                                Dense: -1,
+                                Conv2DTranspose: -1}
+
+    """
+    Map from an activation function name to its min/max output values (if known).
+    The values are used for tensor min/max values initialization.
+    """
+    _activation_min_max_mapping = {SOFTMAX: (0, SOFTMAX_THRESHOLD),
+                                   SIGMOID: (0, 1),
+                                   LINEAR: (None, None),
+                                   IDENTITY: (None, None),
+                                   TANH: (-1, 1),
+                                   SWISH: (-0.279, None),
+                                   RELU: (0, None),
+                                   SELU: (-1.76, None),
+                                   GELU: (-0.17, None),
+                                   }
+
+    """
+    Map from an Keras module to its min/max output values (if known).
+    The values are used for tensor min/max values initialization.
+    """
+    _layer_min_max_mapping = {Softmax: (0, SOFTMAX_THRESHOLD),
+                              ELU: (-1, None),
+                              tf.nn.silu: (-0.279, None),
+                              tf.nn.swish: (-0.279, None),
+                              tf.nn.sigmoid: (0, 1),
+                              tf.nn.tanh: (-1, 1),
+                              tf.nn.relu: (0, None),
+                              tf.nn.relu6: (0, None),
+                              tf.nn.gelu: (-0.17, None),
+                              tf.nn.elu: (-1, None),
+                              tf.nn.selu: (-1.76, None),
+                              tf.nn.softplus: (0, None),
+                              tf.nn.softmax: (0, SOFTMAX_THRESHOLD),
+                              }
+
+    """
+    Mapping from a QuantizationMethod to an activation quantizer function.
+    """
+    activation_quantizer_mapping = {QuantizationMethod.POWER_OF_TWO: power_of_two_quantization,
+                                    QuantizationMethod.SYMMETRIC: symmetric_quantization,
+                                    QuantizationMethod.UNIFORM: uniform_quantization,
+                                    QuantizationMethod.LUT_POT_QUANTIZER: activation_lut_kmean_quantizer}
+
+    @classmethod
+    def get_layer_min_max(cls, layer: Any, fw_attrs: Dict) -> Tuple[float, float]:
+        """
+        Return layer min/max mapping the FrameworkInfo holds.
+        Args:
+            layer: A layer to check if has a min/max known values.
+            fw_attrs: framework attributes from framework layer.
+
+        Returns:
+            Layer's min/max known values.
+        """
+
+        if cls.layers_has_min_max(layer):
+            return cls._layer_min_max_mapping[layer]
+        elif isinstance(layer, Activation) and fw_attrs[ACTIVATION] in cls._activation_min_max_mapping:
+            return cls._activation_min_max_mapping[fw_attrs[ACTIVATION]]
+        else:
+            return None, None
+
+    @classmethod
+    def get_kernel_channels(cls, node_type: Any) -> ChannelAxisMapping:
+        """
+        Returns node's channels mapping from kernel_channels_mapping or framework specific default value.
+        Args:
+            node_type: A node type
+
+        Returns:
+            Node's channels mapping.
+
+        """
+        return cls.kernel_channels_mapping.get(node_type, cls._default_channel_mapping)
+
+    @classmethod
+    def get_out_channel_axis(cls, node_type: Any):
+        """
+        Returns node's output channel mapping from out_channel_axis_mapping or framework specific default value.
+        Args:
+            node_type: A node type.
+
+        Returns:
+            Node's output channel axis.
+
+        """
+        return cls.out_channel_axis_mapping.get(node_type, -1)
+
+
+def set_keras_info(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        set_fw_info(KerasInfo)
+        return func(*args, **kwargs)
+    return wrapper

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

@@ -21,7 +21,6 @@ from model_compression_toolkit.core.common.graph.graph_matchers import NodeOpera
 from model_compression_toolkit.core.common.substitutions.batchnorm_folding import BatchNormalizationFolding, BatchNormalizationForwardFolding
 from model_compression_toolkit.core.keras.constants import KERNEL, LINEAR, ACTIVATION, DEPTHWISE_KERNEL, BIAS, GAMMA, BETA, \
     MOVING_MEAN, MOVING_VARIANCE, EPSILON, USE_BIAS, LAYER_NAME, GROUPS
-from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 
 
 def batchnorm_folding_node_matchers() -> [BaseNode, BaseNode]:
@@ -77,9 +76,7 @@ def update_kernel_for_bn_folding_fn(conv_node: BaseNode,
     else:
         kernel = kernel * weights_scale.reshape((1, 1, 1, -1))
 
-    kernel_name = DEFAULT_KERAS_INFO.get_kernel_op_attributes(conv_node.type)[0]
-
-    return kernel, kernel_name
+    return kernel, conv_node.kernel_attr
 
 
 def update_weights_for_bn_forward_folding_fn(conv_node: BaseNode,
@@ -108,9 +105,7 @@ def update_weights_for_bn_forward_folding_fn(conv_node: BaseNode,
         bias_update = (kernel * bias_factor.reshape((1, 1, -1, 1))).sum(2)
         kernel = kernel * weights_scale.reshape((1, 1, -1, 1))
 
-    kernel_name = DEFAULT_KERAS_INFO.get_kernel_op_attributes(conv_node.type)[0]
-
-    return kernel, bias + bias_update.flatten(), kernel_name
+    return kernel, bias + bias_update.flatten(), conv_node.kernel_attr
 
 
 def get_kernel_hw_fn(kernel: np.ndarray) -> [int, int]:

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

@@ -27,7 +27,6 @@ from model_compression_toolkit.core.common.graph.base_graph import Graph, BaseNo
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
 from model_compression_toolkit.constants import REUSE, REUSE_GROUP
-from model_compression_toolkit.core.keras.reader.node_builder import REUSED_IDENTIFIER
 from model_compression_toolkit.core.keras.constants import KERNEL, BIAS, USE_BIAS, NUM_HEADS, KEY_DIM, VALUE_DIM, \
     QUERY_SHAPE, KEY_SHAPE, VALUE_SHAPE, OUTPUT_SHAPE, ATTENTION_AXES, ACTIVATION, GROUPS, LINEAR, FILTERS, PADDING, \
     FUNCTION, DIMS, TARGET_SHAPE, F_STRIDED_SLICE, F_STACK, Q_KERNEL, Q_BIAS, K_KERNEL, K_BIAS, V_KERNEL, V_BIAS, \

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

@@ -97,16 +97,14 @@ class BaseInputScaling(common.BaseSubstitution):
             scale_factor = threshold_float / threshold
             graph.user_info.set_input_scale(1 / scale_factor)
 
-            kernel_attr = graph.fw_info.get_kernel_op_attributes(linear_layer.type)[0]
-
-            w1_fixed = linear_layer.get_weights_by_keys(kernel_attr) * scale_factor
-            linear_layer.set_weights_by_keys(kernel_attr, w1_fixed)
+            w1_fixed = linear_layer.get_weights_by_keys(linear_layer.kernel_attr) * scale_factor
+            linear_layer.set_weights_by_keys(linear_layer.kernel_attr, w1_fixed)
 
             graph.scale_stats_collector(input_layer, 1 / scale_factor)
 
             # After scaling weights may have different thresholds so it needs to be recalculated
             for nqc in linear_layer.candidates_quantization_cfg:
-                nqc.weights_quantization_cfg.get_attr_config(kernel_attr).calculate_and_set_weights_params(w1_fixed,
+                nqc.weights_quantization_cfg.get_attr_config(linear_layer.kernel_attr).calculate_and_set_weights_params(w1_fixed,
                                                                                                            nqc.weights_quantization_cfg.min_threshold)
 
         return graph

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

@@ -63,17 +63,15 @@ class ScaleEqualization(BaseScaleEqualization):
     """
 
     def __init__(self,
-                 quant_config: QuantizationConfig,
-                 fw_info: FrameworkInfo):
+                 quant_config: QuantizationConfig):
         """
         Initialize a ScaleEqualization object.
         Args:
             quant_config: Quantization configuration.
-            fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
             groups of layers by how they should be quantized, etc.)
         """
 
-        super().__init__(quant_config=quant_config, fw_info=fw_info, matcher_instance=MATCHER,
+        super().__init__(quant_config=quant_config, matcher_instance=MATCHER,
                          kernel_str=KERNEL, bias_str=BIAS)
 
 
@@ -83,17 +81,15 @@ class ScaleEqualizationWithPad(BaseScaleEqualization):
     """
 
     def __init__(self,
-                 quant_config: QuantizationConfig,
-                 fw_info: FrameworkInfo):
+                 quant_config: QuantizationConfig):
         """
         Initialize a ScaleEqualizationWithPad object.
         Args:
             quant_config: Quantization configuration.
-            fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
             groups of layers by how they should be quantized, etc.)
         """
 
-        super().__init__(quant_config=quant_config, fw_info=fw_info, matcher_instance=MATCHER_WITH_PAD,
+        super().__init__(quant_config=quant_config, matcher_instance=MATCHER_WITH_PAD,
                          kernel_str=KERNEL, bias_str=BIAS)
 
 
@@ -104,17 +100,15 @@ class ScaleEqualizationMidActivation(BaseScaleEqualization):
     """
 
     def __init__(self,
-                 quant_config: QuantizationConfig,
-                 fw_info: FrameworkInfo):
+                 quant_config: QuantizationConfig):
         """
         Initialize a ScaleEqualizationMidActivation object.
         Args:
             quant_config: Quantization configuration.
-            fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
             groups of layers by how they should be quantized, etc.)
         """
 
-        super().__init__(quant_config=quant_config, fw_info=fw_info, matcher_instance=MATCHER_MID,
+        super().__init__(quant_config=quant_config, matcher_instance=MATCHER_MID,
                          kernel_str=KERNEL, bias_str=BIAS)
 
 
@@ -124,15 +118,13 @@ class ScaleEqualizationMidActivationWithPad(BaseScaleEqualization):
     """
 
     def __init__(self,
-                 quant_config: QuantizationConfig,
-                 fw_info: FrameworkInfo):
+                 quant_config: QuantizationConfig):
         """
         Initialize a ScaleEqualizationMidActivationWithPad object.
         Args:
             quant_config: Quantization configuration.
-            fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
             groups of layers by how they should be quantized, etc.)
         """
 
-        super().__init__(quant_config=quant_config, fw_info=fw_info, matcher_instance=MATCHER_MID_WITH_PAD,
+        super().__init__(quant_config=quant_config, matcher_instance=MATCHER_MID_WITH_PAD,
                          kernel_str=KERNEL, bias_str=BIAS)