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

model_compression_toolkit/core/common/mixed_precision/sensitivity_eval/metric_calculators.py

@@ -15,7 +15,7 @@
 import numpy as np
 from typing import runtime_checkable, Protocol, Callable, Any, List, Tuple
 
-from model_compression_toolkit.core import FrameworkInfo, MixedPrecisionQuantizationConfig, MpDistanceWeighting
+from model_compression_toolkit.core import MixedPrecisionQuantizationConfig, MpDistanceWeighting
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.hessian import HessianInfoService, HessianScoresRequest, HessianMode, \
     HessianScoresGranularity
@@ -62,15 +62,12 @@ class DistanceMetricCalculator(MetricCalculator):
                  graph: Graph,
                  mp_config: MixedPrecisionQuantizationConfig,
                  representative_data_gen: Callable,
-                 fw_info: FrameworkInfo,
                  fw_impl: Any,
                  hessian_info_service: HessianInfoService = None):
         """
         Args:
             graph: Graph to search for its MP configuration.
             mp_config: MP Quantization configuration for how the graph should be quantized.
-            fw_info: FrameworkInfo object about the specific framework
-                (e.g., attributes of different layers' weights to quantize).
             fw_impl: FrameworkImplementation object with a specific framework methods implementation.
             representative_data_gen: Dataset used for getting batches for inference.
             hessian_info_service: HessianInfoService to fetch Hessian approximation information.
@@ -78,14 +75,13 @@ class DistanceMetricCalculator(MetricCalculator):
         self.graph = graph
         self.mp_config = mp_config
         self.representative_data_gen = representative_data_gen
-        self.fw_info = fw_info
         self.fw_impl = fw_impl
 
         if self.mp_config.distance_weighting_method == MpDistanceWeighting.HESSIAN:
             assert hessian_info_service is not None, ('Expected HessianInfoService object to be passed with Hessian '
                                                       'distance weighting')
 
-        self.sorted_configurable_nodes_names = graph.get_configurable_sorted_nodes_names(self.fw_info)
+        self.sorted_configurable_nodes_names = graph.get_configurable_sorted_nodes_names()
 
         # Get interest points and output points set for distance measurement and set other helper datasets
         # We define a separate set of output nodes of the model for the purpose of sensitivity computation.
@@ -396,8 +392,8 @@ class DistanceMetricCalculator(MetricCalculator):
         """
 
         return [n.node for n in graph.get_outputs()
-                if (graph.fw_info.is_kernel_op(n.node.type) and
-                    n.node.is_weights_quantization_enabled(graph.fw_info.get_kernel_op_attributes(n.node.type)[0])) or
+                if (n.node.is_kernel_op and
+                    n.node.is_weights_quantization_enabled(n.node.kernel_attr)) or
                 n.node.is_activation_quantization_enabled()]
 
     @staticmethod

model_compression_toolkit/core/common/mixed_precision/sensitivity_eval/sensitivity_evaluation.py

@@ -38,7 +38,6 @@ class SensitivityEvaluation:
                  graph: Graph,
                  mp_config: MixedPrecisionQuantizationConfig,
                  representative_data_gen: Callable,
-                 fw_info: FrameworkInfo,
                  fw_impl: Any,
                  disable_activation_for_metric: bool = False,
                  hessian_info_service: HessianInfoService = None
@@ -46,8 +45,6 @@ class SensitivityEvaluation:
         """
         Args:
             graph: Graph to search for its MP configuration.
-            fw_info: FrameworkInfo object about the specific framework
-                (e.g., attributes of different layers' weights to quantize).
             mp_config: MP Quantization configuration for how the graph should be quantized.
             representative_data_gen: Dataset used for getting batches for inference.
             fw_impl: FrameworkImplementation object with a specific framework methods implementation.
@@ -57,14 +54,13 @@ class SensitivityEvaluation:
         """
         self.mp_config = mp_config
         self.representative_data_gen = representative_data_gen
-        self.fw_info = fw_info
         self.fw_impl = fw_impl
 
         if self.mp_config.custom_metric_fn:
             self.metric_calculator = CustomMetricCalculator(graph, self.mp_config.custom_metric_fn)
         else:
             self.metric_calculator = DistanceMetricCalculator(graph, mp_config, representative_data_gen,
-                                                              fw_info=fw_info, fw_impl=fw_impl,
+                                                              fw_impl=fw_impl,
                                                               hessian_info_service=hessian_info_service)
 
         # Build a mixed-precision model which can be configured to use different bitwidth in different layers.
@@ -111,8 +107,7 @@ class SensitivityEvaluation:
 
         model_mp, _, conf_node2layers = self.fw_impl.model_builder(evaluation_graph,
                                                                    mode=ModelBuilderMode.MIXEDPRECISION,
-                                                                   append2output=outputs,
-                                                                   fw_info=self.fw_info)
+                                                                   append2output=outputs)
 
         # Disable all configurable quantizers. They will be activated one at a time during sensitivity evaluation.
         for layer in itertools.chain(*conf_node2layers.values()):

model_compression_toolkit/core/common/model_collector.py

@@ -18,7 +18,7 @@ import numpy as np
 from typing import List, Union, Tuple, Optional
 
 from networkx.algorithms.dag import topological_sort
-from model_compression_toolkit.core import FrameworkInfo, QuantizationErrorMethod
+from model_compression_toolkit.core import QuantizationErrorMethod
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.graph.base_graph import Graph
@@ -30,7 +30,6 @@ from model_compression_toolkit.core.common.collectors.statistics_collector impor
 
 
 def create_stats_collector_for_node(node: common.BaseNode,
-                                    fw_info: FrameworkInfo,
                                     quant_node_in_fln: bool) -> BaseStatsCollector:
     """
     Gets a node and a groups list and create and return a statistics collector for a node
@@ -39,7 +38,7 @@ def create_stats_collector_for_node(node: common.BaseNode,
 
     Args:
         node: Node to create its statistics collector.
-        fw_info: Information relevant to a specific framework about what is out channel axis (for statistics per-channel).
+        quant_node_in_fln: Whether the node should be quantized as part of an FLN.
 
     Returns:
         Statistics collector for statistics collection for the node.
@@ -48,7 +47,7 @@ def create_stats_collector_for_node(node: common.BaseNode,
     if node.is_activation_quantization_enabled() or quant_node_in_fln:
         min_output = getattr(node.prior_info, 'min_output', None)
         max_output = getattr(node.prior_info, 'max_output', None)
-        stats_collector = common.StatsCollector(out_channel_axis=fw_info.out_channel_axis_mapping.get(node.type),
+        stats_collector = common.StatsCollector(out_channel_axis=node.out_channel_axis,
                                                 init_min_value=min_output,
                                                 init_max_value=max_output)
     else:
@@ -58,21 +57,19 @@ def create_stats_collector_for_node(node: common.BaseNode,
 
 
 def create_tensor2node(graph: common.Graph,
-                       node: common.BaseNode,
-                       fw_info: common.FrameworkInfo):
+                       node: common.BaseNode):
     """
     Force statistic collector creation and assignment for a node.
     Args:
         graph: Graph of the node (for retrieving the current tensor).
         node: Node to create a tensor for.
-        fw_info: Specific framework information (for example, output channels index).
 
     """
     current_sc = graph.get_out_stats_collector(node)
     is_list_nostat_collectors = isinstance(current_sc, list) and len(
         [sc for sc in current_sc if not isinstance(sc, common.NoStatsCollector)]) == 0
     if isinstance(current_sc, common.NoStatsCollector) or current_sc is None or is_list_nostat_collectors:
-        stats_collector = common.StatsCollector(fw_info.out_channel_axis_mapping.get(node.type))
+        stats_collector = common.StatsCollector(node.out_channel_axis)
         graph.set_out_stats_collector_to_node(node, stats_collector)
 
 
@@ -140,7 +137,6 @@ class ModelCollector:
 
     def __init__(self, graph: Graph,
                  fw_impl: FrameworkImplementation,
-                 fw_info: FrameworkInfo,
                  hessian_info_service: HessianInfoService = None,
                  qc: common.QuantizationConfig = common.DEFAULTCONFIG):
         """
@@ -149,12 +145,10 @@ class ModelCollector:
         Args:
             graph: Graph to build a model from it.
             fw_impl: FrameworkImplementation object with a specific framework methods implementation.
-            fw_info: FrameworkInfo object with a specific framework information.
             qc: Quantization configuration containing parameters for how the graph should be quantized.
         """
 
         self.fw_impl = fw_impl
-        self.fw_info = fw_info
         self.hessian_service = hessian_info_service
         self.qc = qc
         self.model_outputs = [out.node for out in graph.get_outputs()]
@@ -162,17 +156,15 @@ class ModelCollector:
         # Assign statistics collectors to nodes
         for n in graph.get_topo_sorted_nodes():
             quant_node_in_fln = n.is_fln_quantization() and graph.fusing_info.is_quantized_node_in_fln(n)
-            sc = create_stats_collector_for_node(n, fw_info=fw_info, quant_node_in_fln=quant_node_in_fln)  # Get static collector for the node
+            sc = create_stats_collector_for_node(n, quant_node_in_fln=quant_node_in_fln)  # Get static collector for the node
             # If we use bias correction, and the node has kernel weights to quantize, we need to make sure
             # its previous nodes' tensors are consistent with this node.
-            kernel_attr = fw_info.get_kernel_op_attributes(n.type)[0]
-            if qc.weights_bias_correction and kernel_attr is not None and n.is_weights_quantization_enabled(
-                    kernel_attr):
+            if qc.weights_bias_correction and n.kernel_attr is not None and n.is_weights_quantization_enabled(
+                    n.kernel_attr):
                 for ie in graph.incoming_edges(n):
                     input_node = ie.source_node
                     create_tensor2node(graph,
-                                       input_node,
-                                       fw_info)
+                                       input_node)
             if sc is not None:
                 graph.set_out_stats_collector_to_node(n, sc)
 
@@ -205,13 +197,11 @@ class ModelCollector:
         # TODO: Add integration test for this case
         append2output = outputs_nodes + [n for n in self.model_outputs if n not in outputs_nodes]
 
-
         # Build a float model and output all layers' outputs
         # (that should be collected) as the model's outputs
         self.model, _ = self.fw_impl.model_builder(graph,
                                                    mode=ModelBuilderMode.FLOAT,
-                                                   append2output=append2output,
-                                                   fw_info=self.fw_info)
+                                                   append2output=append2output)
 
     def infer(self, inputs_list: List[np.ndarray]):
         """

model_compression_toolkit/core/common/model_validation.py

@@ -10,17 +10,14 @@ class ModelValidation:
     """
 
     def __init__(self,
-                 model: Any,
-                 fw_info:FrameworkInfo):
+                 model: Any):
         """
         Initialize a ModelValidation object.
 
         Args:
             model: Model to check its validity.
-            fw_info: Information about the specific framework of the model.
         """
         self.model = model
-        self.fw_info = fw_info
 
     @abstractmethod
     def validate_output_channel_consistency(self):

model_compression_toolkit/core/common/network_editors/actions.py

@@ -22,7 +22,7 @@ from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.logger import Logger
 
 
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+from model_compression_toolkit.core.common.framework_info import get_fw_info
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
 from model_compression_toolkit.core.common.quantization.quantization_params_fn_selection import \
     get_activation_quantization_params_fn, get_weights_quantization_params_fn
@@ -64,15 +64,13 @@ class BaseAction(ABC):
     """
 
     @abstractmethod
-    def apply(self, node: BaseNode, graph, fw_info):
+    def apply(self, node: BaseNode, graph):
         """
         Apply an action on the node after matching the node with a node filter.
 
         Args:
             node: Node to apply the action on.
             graph: Graph to apply the action on.
-            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.)
 
         Returns:
             Node after action is applied.
@@ -95,15 +93,13 @@ class ChangeCandidatesWeightsQuantConfigAttr(BaseAction):
         self.kwargs = kwargs
         self.attr_name = attr_name
 
-    def apply(self, node: BaseNode, graph, fw_info):
+    def apply(self, node: BaseNode, graph):
         """
         Change the attribute 'attr_name' in weights quantization config candidates with 'attr_value'.
 
         Args:
             node: Node object to change its quant_config.
             graph: Graph to apply the action on.
-            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.)
         Returns:
             The node after its weights' quantization config candidates have been modified.
         """
@@ -128,7 +124,7 @@ class ChangeFinalWeightsQuantConfigAttr(BaseAction):
         self.kwargs = kwargs
         self.attr_name = attr_name
 
-    def apply(self, node: BaseNode, graph, fw_info):
+    def apply(self, node: BaseNode, graph):
         if node.final_weights_quantization_cfg is not None:
             for parameter_name, parameter_value in self.kwargs.items():
                 node.final_weights_quantization_cfg.set_quant_config_attr(parameter_name, parameter_value,
@@ -147,17 +143,13 @@ class ChangeCandidatesActivationQuantConfigAttr(BaseAction):
         """
         self.kwargs = kwargs
 
-    def apply(self, node: BaseNode, graph, fw_info):
+    def apply(self, node: BaseNode, graph):
         """
         Change the attribute 'attr_name' in activation quantization configuration candidates with 'attr_value'.
 
         Args:
             node: Node object to change its quant_config.
             graph: Graph to apply the action on.
-            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.)
-        Returns:q
-            The node after its activation quantization configuration candidates have been modified.
         """
         for nqc in node.candidates_quantization_cfg:
             for parameter_name, parameter_value in self.kwargs.items():
@@ -176,7 +168,7 @@ class ChangeFinalActivationQuantConfigAttr(BaseAction):
         """
         self.kwargs = kwargs
 
-    def apply(self, node: BaseNode, graph, fw_info):
+    def apply(self, node: BaseNode, graph):
         if node.final_activation_quantization_cfg is not None:
             for parameter_name, parameter_value in self.kwargs.items():
                 node.final_activation_quantization_cfg.set_quant_config_attr(parameter_name, parameter_value)
@@ -203,15 +195,13 @@ class ChangeQuantizationParamFunction(BaseAction):
         self.weights_quantization_params_fn = weights_quantization_params_fn
         self.attr_name = attr_name
 
-    def apply(self, node: BaseNode, graph, fw_info):
+    def apply(self, node: BaseNode, graph):
         """
         Change the node's weights/activations quantization params function.
 
         Args:
             node: Node object to change its quantization params function.
             graph: Graph to apply the action on.
-            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.)
 
         Returns:
             The node after its quantization params function has been modified.
@@ -240,15 +230,13 @@ class ChangeFinalActivationQuantizationMethod(BaseAction):
 
         self.activation_quantization_method = activation_quantization_method
 
-    def apply(self, node: BaseNode, graph, fw_info):
+    def apply(self, node: BaseNode, graph):
         """
         Change the node's activations quantization function.
 
         Args:
             node: Node object to change its threshold selection function.
             graph: Graph to apply the action on.
-            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.)
 
         Returns:
             The node after its quantization function has been modified.
@@ -262,7 +250,7 @@ class ChangeFinalActivationQuantizationMethod(BaseAction):
             node.final_activation_quantization_cfg.set_activation_quantization_params_fn(
                 activation_quantization_params_fn)
 
-            activation_quantization_fn = fw_info.activation_quantizer_mapping.get(self.activation_quantization_method)
+            activation_quantization_fn = get_fw_info().activation_quantizer_mapping.get(self.activation_quantization_method)
 
             node.final_activation_quantization_cfg.set_activation_quantization_fn(activation_quantization_fn)
             node.final_activation_quantization_cfg.activation_quantization_method = self.activation_quantization_method
@@ -282,18 +270,14 @@ class ChangeCandidatesActivationQuantizationMethod(BaseAction):
         """
         self.activation_quantization_method = activation_quantization_method
 
-    def apply(self, node: BaseNode, graph, fw_info):
+    def apply(self, node: BaseNode, graph):
         """
         Change the node's activations quantization function.
 
         Args:
             node: Node object to change its threshold selection function.
             graph: Graph to apply the action on.
-            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.)
 
-        Returns:
-            The node after its quantization function has been modified.
         """
         if self.activation_quantization_method is not None:
             for qc in node.candidates_quantization_cfg:
@@ -301,7 +285,7 @@ class ChangeCandidatesActivationQuantizationMethod(BaseAction):
                     self.activation_quantization_method)
 
                 qc.activation_quantization_cfg.set_activation_quantization_params_fn(activation_quantization_params_fn)
-                activation_quantization_fn = fw_info.activation_quantizer_mapping.get(
+                activation_quantization_fn = get_fw_info().activation_quantizer_mapping.get(
                     self.activation_quantization_method)
 
                 if activation_quantization_fn is None:
@@ -328,18 +312,14 @@ class ChangeFinalWeightsQuantizationMethod(BaseAction):
         self.weights_quantization_method = weights_quantization_method
         self.attr_name = attr_name
 
-    def apply(self, node: BaseNode, graph, fw_info):
+    def apply(self, node: BaseNode, graph):
         """
         Change the node's weights quantization function.
 
         Args:
             node: Node object to change its threshold selection function.
             graph: Graph to apply the action on.
-            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.)
 
-        Returns:
-            The node after its quantization function has been modified.
         """
 
         if self.weights_quantization_method is not None and node.final_weights_quantization_cfg is not None:
@@ -376,15 +356,13 @@ class ChangeCandidatesWeightsQuantizationMethod(BaseAction):
         self.weights_quantization_method = weights_quantization_method
         self.attr_name = attr_name
 
-    def apply(self, node: BaseNode, graph: Graph, fw_info: FrameworkInfo):
+    def apply(self, node: BaseNode, graph: Graph):
         """
         Change the node's weights quantization function.
 
         Args:
             node: Node object to change its threshold selection function.
             graph: Graph to apply the action on.
-            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.)
 
         Returns:
             The node after its quantization function has been modified.
@@ -422,15 +400,13 @@ class ReplaceLayer(BaseAction):
         self.layer_type = layer_type
         self.get_params_and_weights_fn = get_params_and_weights_fn
 
-    def apply(self, node: BaseNode, graph: Graph, fw_info: FrameworkInfo):
+    def apply(self, node: BaseNode, graph: Graph):
         """
         Replacing node's layer type and configurations
 
         Args:
             node: Node object to replace or modify
             graph: Graph to apply the action on.
-            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.)
 
         Returns:
             The node after its layer functionality has been modified.

model_compression_toolkit/core/common/network_editors/edit_network.py

@@ -14,20 +14,17 @@
 # ==============================================================================
 from typing import List
 
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.network_editors import EditRule
 
 
 def edit_network_graph(graph: Graph,
-                       fw_info: FrameworkInfo,
                        network_editor: List[EditRule]):
     """
     Apply a list of edit rules on a graph.
 
     Args:
         graph: The graph to edit.
-        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.)
         network_editor: List of edit rules to apply to the graph.
 
@@ -38,5 +35,5 @@ def edit_network_graph(graph: Graph,
     for edit_rule in network_editor:
         filtered_nodes = graph.filter(edit_rule.filter)
         for node in filtered_nodes:
-            edit_rule.action.apply(node, graph, fw_info)
+            edit_rule.action.apply(node, graph)
     # return graph

model_compression_toolkit/core/common/pruning/channels_grouping.py

@@ -26,18 +26,14 @@ class ChannelGrouping:
     based on their importance scores and SIMD group sizes.
     """
 
-    def __init__(self,
-                 prunable_nodes: List[BaseNode],
-                 fw_info: FrameworkInfo):
+    def __init__(self, prunable_nodes: List[BaseNode]):
         """
         Initializes the ChannelGrouping with necessary information.
 
         Args:
             prunable_nodes: List of nodes that can be pruned.
-            fw_info: Framework-specific information and utilities.
         """
         self.prunable_nodes = prunable_nodes
-        self.fw_info = fw_info
         # Store for each node a list of numpy arrays. Each numpy array represents the
         # indices of the channels in an SIMD group.
         self._simd_groups_indices = {}

model_compression_toolkit/core/common/pruning/greedy_mask_calculator.py

@@ -38,7 +38,6 @@ class GreedyMaskCalculator:
     """
     def __init__(self,
                  prunable_nodes: List[BaseNode],
-                 fw_info: FrameworkInfo,
                  simd_groups_scores: Dict[BaseNode, np.ndarray],
                  target_resource_utilization: ResourceUtilization,
                  graph: Graph,
@@ -48,7 +47,6 @@ class GreedyMaskCalculator:
         """
         Args:
             prunable_nodes (List[BaseNode]): Nodes that are eligible for pruning.
-            fw_info (FrameworkInfo): Framework-specific information and utilities.
             simd_groups_scores (Dict[BaseNode, np.ndarray]): Importance scores for each SIMG group in a prunable node.
             target_resource_utilization (ResourceUtilization): The target resource utilization to achieve.
             graph (Graph): The computational graph of the model.
@@ -57,7 +55,6 @@ class GreedyMaskCalculator:
             simd_groups_indices (Dict[BaseNode, List[List[int]]]): Indices of SIMD groups in each node.
         """
         self.prunable_nodes = prunable_nodes
-        self.fw_info = fw_info
         self.target_resource_utilization = target_resource_utilization
         self.graph = graph
         self.fw_impl = fw_impl
@@ -67,14 +64,11 @@ class GreedyMaskCalculator:
         self.simd_groups_scores = simd_groups_scores
 
         self.oc_pruning_mask = PerSIMDGroupMask(prunable_nodes=prunable_nodes,
-                                                fw_info=fw_info,
                                                 simd_groups_indices=simd_groups_indices)
 
         self.memory_calculator = MemoryCalculator(graph=graph,
-                                                  fw_info=fw_info,
                                                   fw_impl=fw_impl)
 
-
     def get_mask(self) -> Dict[BaseNode, np.ndarray]:
         """
         Retrieves the current pruning mask for each prunable node.

model_compression_toolkit/core/common/pruning/importance_metrics/lfh_importance_metric.py

@@ -38,8 +38,7 @@ class LFHImportanceMetric(BaseImportanceMetric):
                  graph: Graph,
                  representative_data_gen: Callable,
                  fw_impl: PruningFrameworkImplementation,
-                 pruning_config: PruningConfig,
-                 fw_info: FrameworkInfo):
+                 pruning_config: PruningConfig):
         """
         Initialize the LFHImportanceMetric instance.
 
@@ -48,13 +47,11 @@ class LFHImportanceMetric(BaseImportanceMetric):
             representative_data_gen (Callable): Function to generate representative data.
             fw_impl (PruningFrameworkImplementation): Implementation of pruning for the framework.
             pruning_config (PruningConfig): Configuration for pruning.
-            fw_info (FrameworkInfo): Framework-specific information.
         """
         self.float_graph = graph
         self.representative_data_gen = representative_data_gen
         self.fw_impl = fw_impl
         self.pruning_config = pruning_config
-        self.fw_info = fw_info
 
         # Initialize internal dictionaries for storing intermediate computations.
         self._entry_node_to_hessian_score = {}
@@ -158,8 +155,7 @@ class LFHImportanceMetric(BaseImportanceMetric):
             Dict[BaseNode, List[np.ndarray]]: Dictionary of entry nodes mapped to their SIMD group indices.
         """
         # Initialize channel grouping utility.
-        channel_grouping = ChannelGrouping(prunable_nodes=list(entry_node_to_score.keys()),
-                                           fw_info=self.fw_info)
+        channel_grouping = ChannelGrouping(prunable_nodes=list(entry_node_to_score.keys()))
 
         channel_grouping.group_scores_by_simd_groups(entry_node_to_score)
         grouped_indices = channel_grouping.simd_groups_indices
@@ -249,20 +245,14 @@ class LFHImportanceMetric(BaseImportanceMetric):
         Returns:
             tuple: A tuple containing the kernel attribute, the number of output channels, and the axis of the output channels.
         """
-        kernel_attr = self.fw_info.get_kernel_op_attributes(entry_node.type)
-        # Ensure only one kernel attribute exists for the given node.
-        if len(kernel_attr) != 1:
-            Logger.critical(f"Expected a single attribute but found multiple attributes ({len(kernel_attr)}) for node {entry_node}.")
-        kernel_attr = kernel_attr[0]
-
         # Retrieve and validate the axis index for the output channels.
-        oc_axis, _ = self.fw_info.kernel_channels_mapping.get(entry_node.type)
+        oc_axis = entry_node.channel_axis.output
         if oc_axis is None or int(oc_axis) != oc_axis:
             Logger.critical(f"Invalid output channel axis type for node {entry_node}: expected integer but got {oc_axis}.")
 
         # Get the number of output channels based on the kernel attribute and axis.
-        num_oc = entry_node.get_weights_by_keys(kernel_attr[0]).shape[oc_axis]
-        return kernel_attr, num_oc, oc_axis
+        num_oc = entry_node.get_weights_by_keys(entry_node.kernel_attr).shape[oc_axis]
+        return entry_node.kernel_attr, num_oc, oc_axis
 
     def _concatenate_tensors_by_indices(self,
                                         channels: List[np.ndarray],

model_compression_toolkit/core/common/pruning/mask/per_channel_mask.py

@@ -35,9 +35,8 @@ class MaskIndicator(Enum):
     REMAINED = 1
 
 
-
 class PerChannelMask:
-    def __init__(self, prunable_nodes: List[BaseNode], fw_info: FrameworkInfo):
+    def __init__(self, prunable_nodes: List[BaseNode]):
         """
         Initializes the PerChannelMask with prunable nodes and framework information.
         This class is responsible for maintaining and updating the pruning masks for each
@@ -46,10 +45,8 @@ class PerChannelMask:
 
         Args:
             prunable_nodes: List of nodes in the model that are subject to pruning.
-            fw_info: Framework-specific information required for pruning operations.
         """
         self.prunable_nodes = prunable_nodes
-        self.fw_info = fw_info
         self._mask = None  # Initialize the mask dictionary
         self._init_masks()  # Call to initialize masks for each prunable node
 
@@ -106,8 +103,7 @@ class PerChannelMask:
         Returns:
             int: Number of output channels for the node.
         """
-        kernel_attr = self.fw_info.get_kernel_op_attributes(node.type)[0]
-        oc_axis = self.fw_info.kernel_channels_mapping.get(node.type)[0]
-        num_oc = node.get_weights_by_keys(kernel_attr).shape[oc_axis]
+        oc_axis = node.channel_axis.output
+        num_oc = node.get_weights_by_keys(node.kernel_attr).shape[oc_axis]
         return num_oc
 

model_compression_toolkit/core/common/pruning/mask/per_simd_group_mask.py

@@ -24,10 +24,10 @@ from model_compression_toolkit.core.common.pruning.memory_calculator import Memo
 from model_compression_toolkit.core.common.pruning.pruning_framework_implementation import PruningFrameworkImplementation
 from model_compression_toolkit.logger import Logger
 
+
 class PerSIMDGroupMask:
     def __init__(self,
                  prunable_nodes: List[BaseNode],
-                 fw_info: FrameworkInfo,
                  simd_groups_indices: Dict[BaseNode, List[List[int]]]):
         """
         Initializes a mask calculator for SIMD groups in prunable nodes.
@@ -35,13 +35,11 @@ class PerSIMDGroupMask:
 
         Args:
             prunable_nodes: List of nodes that can be pruned.
-            fw_info: Framework-specific information.
             simd_groups_indices: A dictionary mapping each node to its SIMD groups' indices.
         """
         # Initialize the per-channel mask
-        self.per_channel_mask = PerChannelMask(prunable_nodes=prunable_nodes, fw_info=fw_info)
+        self.per_channel_mask = PerChannelMask(prunable_nodes=prunable_nodes)
         self.prunable_nodes = prunable_nodes
-        self.fw_info = fw_info
         self.simd_groups_indices = simd_groups_indices
         self._mask_simd = None  # Initialize the SIMD group mask dictionary
         self._init_masks()  # Initialize masks for each prunable node