mct-nightly 2.2.0.20241202.131715__tar.gz → 2.2.0.20241204.524__tar.gz

{mct-nightly-2.2.0.20241202.131715 → mct-nightly-2.2.0.20241204.524}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.2.0.20241202.131715
+Version: 2.2.0.20241204.524
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct-nightly-2.2.0.20241202.131715 → mct-nightly-2.2.0.20241204.524}/mct_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.2.0.20241202.131715
+Version: 2.2.0.20241204.524
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct-nightly-2.2.0.20241202.131715 → mct-nightly-2.2.0.20241204.524}/model_compression_toolkit/__init__.py

@@ -27,4 +27,4 @@ from model_compression_toolkit import data_generation
 from model_compression_toolkit import pruning
 from model_compression_toolkit.trainable_infrastructure.keras.load_model import keras_load_quantized_model
 
-__version__ = "2.2.0.20241202.131715"
+__version__ = "2.2.0.20241204.000524"

{mct-nightly-2.2.0.20241202.131715 → mct-nightly-2.2.0.20241204.524}/model_compression_toolkit/core/common/graph/base_graph.py

@@ -542,9 +542,14 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         """
         # configurability is only relevant for kernel attribute quantization
         potential_conf_nodes = [n for n in list(self) if fw_info.is_kernel_op(n.type)]
-        return list(filter(lambda n: n.is_weights_quantization_enabled(fw_info.get_kernel_op_attributes(n.type)[0])
-                                     and not n.is_all_weights_candidates_equal(fw_info.get_kernel_op_attributes(n.type)[0])
-                                     and (not n.reuse or include_reused_nodes), potential_conf_nodes))
+
+        def is_configurable(n):
+            kernel_attr = fw_info.get_kernel_op_attributes(n.type)[0]
+            return (n.is_weights_quantization_enabled(kernel_attr) and
+                    not n.is_all_weights_candidates_equal(kernel_attr) and
+                    (not n.reuse or include_reused_nodes))
+
+        return [n for n in potential_conf_nodes if is_configurable(n)]
 
     def get_sorted_weights_configurable_nodes(self,
                                               fw_info: FrameworkInfo,
@@ -571,8 +576,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         Returns:
             A list of nodes that their activation can be configured (namely, has one or more activation qc candidate).
         """
-        return list(filter(lambda n: n.is_activation_quantization_enabled()
-                                     and not n.is_all_activation_candidates_equal(), list(self)))
+        return [n for n in list(self) if n.is_activation_quantization_enabled() and not n.is_all_activation_candidates_equal()]
 
     def get_sorted_activation_configurable_nodes(self) -> List[BaseNode]:
         """

{mct-nightly-2.2.0.20241202.131715 → mct-nightly-2.2.0.20241204.524}/model_compression_toolkit/core/common/graph/base_node.py

@@ -384,9 +384,8 @@ class BaseNode:
         """
         # note that if the given attribute name does not exist in the node's attributes mapping,
         # the inner method would log an exception.
-        return all(attr_candidate ==
-                   self.candidates_quantization_cfg[0].weights_quantization_cfg.get_attr_config(attr)
-                   for attr_candidate in self.get_all_weights_attr_candidates(attr))
+        candidates = self.get_all_weights_attr_candidates(attr)
+        return all(candidate == candidates[0] for candidate in candidates[1:])
 
     def has_kernel_weight_to_quantize(self, fw_info):
         """

mct-nightly-2.2.0.20241204.524/model_compression_toolkit/core/common/mixed_precision/mixed_precision_candidates_filter.py

@@ -0,0 +1,77 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import numpy as np
+
+from model_compression_toolkit.core import ResourceUtilization, FrameworkInfo
+from model_compression_toolkit.core.common import Graph
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
+
+
+def filter_candidates_for_mixed_precision(graph: Graph,
+                                          target_resource_utilization: ResourceUtilization,
+                                          fw_info: FrameworkInfo,
+                                          tpc: TargetPlatformCapabilities):
+    """
+    Filters out candidates in case of mixed precision search for only weights or activation compression.
+    For instance, if running only weights compression - filters out candidates of activation configurable nodes
+    such that only a single candidate would remain, with the bitwidth equal to the one defined in the matching layer's
+    base config in the TPC.
+
+    Note" This function modifies the graph inplace!
+
+    Args:
+        graph: A graph representation of the model to be quantized.
+        target_resource_utilization: The resource utilization of the target device.
+        fw_info: fw_info: Information needed for quantization about the specific framework.
+        tpc: TargetPlatformCapabilities object that describes the desired inference target platform.
+
+    """
+
+    tru = target_resource_utilization
+    if tru.total_mem_restricted() or tru.bops_restricted():
+        return
+
+    if tru.weight_restricted() and not tru.activation_restricted():
+        # Running mixed precision for weights compression only -
+        # filter out candidates activation only configurable node
+        weights_conf = graph.get_weights_configurable_nodes(fw_info)
+        activation_configurable_nodes = [n for n in graph.get_activation_configurable_nodes() if n not in weights_conf]
+        for n in activation_configurable_nodes:
+            base_cfg_nbits = n.get_qco(tpc).base_config.activation_n_bits
+            filtered_conf = [c for c in n.candidates_quantization_cfg if
+                             c.activation_quantization_cfg.enable_activation_quantization and
+                             c.activation_quantization_cfg.activation_n_bits == base_cfg_nbits]
+
+            if len(filtered_conf) != 1:
+                Logger.critical(f"Running weights only mixed precision failed on layer {n.name} with multiple "
+                                f"activation quantization configurations.")  # pragma: no cover
+            n.candidates_quantization_cfg = filtered_conf
+
+    elif tru.activation_restricted() and not tru.weight_restricted():
+        # Running mixed precision for activation compression only -
+        # filter out candidates weights only configurable node
+        activation_conf = graph.get_activation_configurable_nodes()
+        weight_configurable_nodes = [n for n in graph.get_weights_configurable_nodes(fw_info) if n not in activation_conf]
+        for n in weight_configurable_nodes:
+            kernel_attr = graph.fw_info.get_kernel_op_attributes(n.type)[0]
+            base_cfg_nbits = n.get_qco(tpc).base_config.attr_weights_configs_mapping[kernel_attr].weights_n_bits
+            filtered_conf = [c for c in n.candidates_quantization_cfg if
+                             c.weights_quantization_cfg.get_attr_config(kernel_attr).enable_weights_quantization and
+                             c.weights_quantization_cfg.get_attr_config(kernel_attr).weights_n_bits == base_cfg_nbits]
+            if len(filtered_conf) != 1:
+                Logger.critical(f"Running activation only mixed precision failed on layer {n.name} with multiple "
+                                f"weights quantization configurations.")  # pragma: no cover
+            n.candidates_quantization_cfg = filtered_conf

{mct-nightly-2.2.0.20241202.131715 → mct-nightly-2.2.0.20241204.524}/model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py

@@ -83,16 +83,17 @@ def search_bit_width(graph_to_search_cfg: Graph,
 
     # Set graph for MP search
     graph = copy.deepcopy(graph_to_search_cfg)  # Copy graph before searching
-    if target_resource_utilization.bops < np.inf:
+    if target_resource_utilization.bops_restricted():
         # Since Bit-operations count target resource utilization is set, we need to reconstruct the graph for the MP search
         graph = substitute(graph, fw_impl.get_substitutions_virtual_weights_activation_coupling())
 
     # If we only run weights compression with MP than no need to consider activation quantization when computing the
     # MP metric (it adds noise to the computation)
-    disable_activation_for_metric = (target_resource_utilization.weights_memory < np.inf and
-                                    (target_resource_utilization.activation_memory == np.inf and
-                                     target_resource_utilization.total_memory == np.inf and
-                                     target_resource_utilization.bops == np.inf)) or graph_to_search_cfg.is_single_activation_cfg()
+    tru = target_resource_utilization
+    weight_only_restricted = tru.weight_restricted() and not (tru.activation_restricted() or
+                                                              tru.total_mem_restricted() or
+                                                              tru.bops_restricted())
+    disable_activation_for_metric = weight_only_restricted or graph_to_search_cfg.is_single_activation_cfg()
 
     # Set Sensitivity Evaluator for MP search. It should always work with the original MP graph,
     # even if a virtual graph was created (and is used only for BOPS utilization computation purposes)
@@ -117,11 +118,10 @@ def search_bit_width(graph_to_search_cfg: Graph,
                                                  target_resource_utilization,
                                                  original_graph=graph_to_search_cfg)
 
-    if search_method in search_methods:  # Get a specific search function
-        search_method_fn = search_methods.get(search_method)
-    else:
-        raise NotImplemented  # pragma: no cover
+    if search_method not in search_methods:
+        raise NotImplementedError()  # pragma: no cover
 
+    search_method_fn = search_methods[search_method]
     # Search for the desired mixed-precision configuration
     result_bit_cfg = search_method_fn(search_manager,
                                       target_resource_utilization)

{mct-nightly-2.2.0.20241202.131715 → mct-nightly-2.2.0.20241204.524}/model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py

@@ -211,14 +211,8 @@ class MixedPrecisionSearchManager:
         Returns: Node's resource utilization vector.
 
         """
-        return self.compute_ru_functions[target][0](
-            self.replace_config_in_index(
-                self.min_ru_config,
-                conf_node_idx,
-                candidate_idx),
-            self.graph,
-            self.fw_info,
-            self.fw_impl)
+        cfg = self.replace_config_in_index(self.min_ru_config, conf_node_idx, candidate_idx)
+        return self.compute_ru_functions[target].metric_fn(cfg, self.graph, self.fw_info, self.fw_impl)
 
     @staticmethod
     def replace_config_in_index(mp_cfg: List[int], idx: int, value: int) -> List[int]:
@@ -253,7 +247,7 @@ class MixedPrecisionSearchManager:
             if target == RUTarget.BOPS:
                 ru_vector = None
             else:
-                ru_vector = self.compute_ru_functions[target][0]([], self.graph, self.fw_info, self.fw_impl)
+                ru_vector = self.compute_ru_functions[target].metric_fn([], self.graph, self.fw_info, self.fw_impl)
 
             non_conf_ru_dict[target] = ru_vector
 
@@ -282,9 +276,9 @@ class MixedPrecisionSearchManager:
                 configurable_nodes_ru_vector = ru_fns[0](config, self.original_graph, self.fw_info, self.fw_impl)
             non_configurable_nodes_ru_vector = self.non_conf_ru_dict.get(ru_target)
             if non_configurable_nodes_ru_vector is None or len(non_configurable_nodes_ru_vector) == 0:
-                ru_ru = self.compute_ru_functions[ru_target][1](configurable_nodes_ru_vector, False)
+                ru_ru = self.compute_ru_functions[ru_target].aggregate_fn(configurable_nodes_ru_vector, False)
             else:
-                ru_ru = self.compute_ru_functions[ru_target][1](
+                ru_ru = self.compute_ru_functions[ru_target].aggregate_fn(
                     np.concatenate([configurable_nodes_ru_vector, non_configurable_nodes_ru_vector]), False)
 
             ru_dict[ru_target] = ru_ru[0]

{mct-nightly-2.2.0.20241202.131715 → mct-nightly-2.2.0.20241204.524}/model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization.py

@@ -72,6 +72,18 @@ class ResourceUtilization:
                f"Total_memory: {self.total_memory}, " \
                f"BOPS: {self.bops}"
 
+    def weight_restricted(self):
+        return self.weights_memory < np.inf
+
+    def activation_restricted(self):
+        return self.activation_memory < np.inf
+
+    def total_mem_restricted(self):
+        return self.total_memory < np.inf
+
+    def bops_restricted(self):
+        return self.bops < np.inf
+
     def get_resource_utilization_dict(self) -> Dict[RUTarget, float]:
         """
         Returns: a dictionary with the ResourceUtilization object's values for each resource utilization target.

{mct-nightly-2.2.0.20241202.131715 → mct-nightly-2.2.0.20241204.524}/model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/ru_functions_mapping.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+from typing import NamedTuple
+
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import RUTarget
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.ru_aggregation_methods import MpRuAggregation
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.ru_methods import MpRuMetric
@@ -20,7 +22,12 @@ from model_compression_toolkit.core.common.mixed_precision.resource_utilization_
 # When adding a RUTarget that we want to consider in our mp search,
 # a matching pair of resource_utilization_tools computation function and a resource_utilization_tools
 # aggregation function should be added to this dictionary
-ru_functions_mapping = {RUTarget.WEIGHTS: (MpRuMetric.WEIGHTS_SIZE, MpRuAggregation.SUM),
-                        RUTarget.ACTIVATION: (MpRuMetric.ACTIVATION_OUTPUT_SIZE, MpRuAggregation.MAX),
-                        RUTarget.TOTAL: (MpRuMetric.TOTAL_WEIGHTS_ACTIVATION_SIZE, MpRuAggregation.TOTAL),
-                        RUTarget.BOPS: (MpRuMetric.BOPS_COUNT, MpRuAggregation.SUM)}
+class RuFunctions(NamedTuple):
+    metric_fn: MpRuMetric
+    aggregate_fn: MpRuAggregation
+
+
+ru_functions_mapping = {RUTarget.WEIGHTS: RuFunctions(MpRuMetric.WEIGHTS_SIZE, MpRuAggregation.SUM),
+                        RUTarget.ACTIVATION: RuFunctions(MpRuMetric.ACTIVATION_OUTPUT_SIZE, MpRuAggregation.MAX),
+                        RUTarget.TOTAL: RuFunctions(MpRuMetric.TOTAL_WEIGHTS_ACTIVATION_SIZE, MpRuAggregation.TOTAL),
+                        RUTarget.BOPS: RuFunctions(MpRuMetric.BOPS_COUNT, MpRuAggregation.SUM)}

{mct-nightly-2.2.0.20241202.131715 → mct-nightly-2.2.0.20241204.524}/model_compression_toolkit/core/common/mixed_precision/search_methods/linear_programming.py

@@ -221,9 +221,9 @@ def _add_set_of_ru_constraints(search_manager: MixedPrecisionSearchManager,
     # search_manager.compute_ru_functions contains a pair of ru_metric and ru_aggregation for each ru target
     # get aggregated ru, considering both configurable and non-configurable nodes
     if non_conf_ru_vector is None or len(non_conf_ru_vector) == 0:
-        aggr_ru = search_manager.compute_ru_functions[target][1](ru_sum_vector)
+        aggr_ru = search_manager.compute_ru_functions[target].aggregate_fn(ru_sum_vector)
     else:
-        aggr_ru = search_manager.compute_ru_functions[target][1](np.concatenate([ru_sum_vector, non_conf_ru_vector]))
+        aggr_ru = search_manager.compute_ru_functions[target].aggregate_fn(np.concatenate([ru_sum_vector, non_conf_ru_vector]))
 
     for v in aggr_ru:
         if isinstance(v, float):
@@ -261,9 +261,7 @@ def _build_layer_to_metrics_mapping(search_manager: MixedPrecisionSearchManager,
     Logger.info('Starting to evaluate metrics')
     layer_to_metrics_mapping = {}
 
-    is_bops_target_resource_utilization = target_resource_utilization.bops < np.inf
-
-    if is_bops_target_resource_utilization:
+    if target_resource_utilization.bops_restricted():
         origin_max_config = search_manager.config_reconstruction_helper.reconstruct_config_from_virtual_graph(search_manager.max_ru_config)
         max_config_value = search_manager.compute_metric_fn(origin_max_config)
     else:
@@ -284,7 +282,7 @@ def _build_layer_to_metrics_mapping(search_manager: MixedPrecisionSearchManager,
             mp_model_configuration[node_idx] = bitwidth_idx
 
             # Build a distance matrix using the function we got from the framework implementation.
-            if is_bops_target_resource_utilization:
+            if target_resource_utilization.bops_restricted():
                 # Reconstructing original graph's configuration from virtual graph's configuration
                 origin_mp_model_configuration = \
                     search_manager.config_reconstruction_helper.reconstruct_config_from_virtual_graph(

{mct-nightly-2.2.0.20241202.131715 → mct-nightly-2.2.0.20241204.524}/model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py

@@ -90,7 +90,7 @@ class SensitivityEvaluation:
                                                       quant_config.num_interest_points_factor)
 
         # We use normalized MSE when not running hessian-based. For Hessian-based normalized MSE is not needed
-        # beacause hessian weights already do normalization.
+        # because hessian weights already do normalization.
         use_normalized_mse = self.quant_config.use_hessian_based_scores is False
         self.ips_distance_fns, self.ips_axis = self._init_metric_points_lists(self.interest_points, use_normalized_mse)
 
@@ -116,14 +116,11 @@ class SensitivityEvaluation:
         # Build images batches for inference comparison
         self.images_batches = self._get_images_batches(quant_config.num_of_images)
 
-        # Get baseline model inference on all samples
-        self.baseline_tensors_list = []  # setting from outside scope
-
         # Casting images tensors to the framework tensor type.
-        self.images_batches = list(map(lambda in_arr: self.fw_impl.to_tensor(in_arr), self.images_batches))
+        self.images_batches = [self.fw_impl.to_tensor(img) for img in self.images_batches]
 
         # Initiating baseline_tensors_list since it is not initiated in SensitivityEvaluationManager init.
-        self._init_baseline_tensors_list()
+        self.baseline_tensors_list = self._init_baseline_tensors_list()
 
         # Computing Hessian-based scores for weighted average distance metric computation (only if requested),
         # and assigning distance_weighting method accordingly.
@@ -193,11 +190,9 @@ class SensitivityEvaluation:
 
     def _init_baseline_tensors_list(self):
         """
-        Evaluates the baseline model on all images and saves the obtained lists of tensors in a list for later use.
-        Initiates a class variable self.baseline_tensors_list
+        Evaluates the baseline model on all images and returns the obtained lists of tensors in a list for later use.
         """
-        self.baseline_tensors_list = [self.fw_impl.to_numpy(self.fw_impl.sensitivity_eval_inference(self.baseline_model,
-                                                                                                    images))
+        return [self.fw_impl.to_numpy(self.fw_impl.sensitivity_eval_inference(self.baseline_model, images))
                                       for images in self.images_batches]
 
     def _build_models(self) -> Any:
@@ -454,7 +449,7 @@ def get_mp_interest_points(graph: Graph,
 
     """
     sorted_nodes = graph.get_topo_sorted_nodes()
-    ip_nodes = list(filter(lambda n: interest_points_classifier(n), sorted_nodes))
+    ip_nodes = [n for n in sorted_nodes if interest_points_classifier(n)]
 
     interest_points_nodes = bound_num_interest_points(ip_nodes, num_ip_factor)
 

{mct-nightly-2.2.0.20241202.131715 → mct-nightly-2.2.0.20241204.524}/model_compression_toolkit/core/keras/back2framework/mixed_precision_model_builder.py

@@ -92,11 +92,8 @@ class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
         if kernel_attr is not None and n.is_weights_quantization_enabled(kernel_attr):
             weights_conf_nodes_names = [node.name for node in self.graph.get_weights_configurable_nodes(self.fw_info)]
             if n.name in weights_conf_nodes_names:
-                return KerasQuantizationWrapper(layer,
-                                                weights_quantizers={
-                                                    kernel_attr: ConfigurableWeightsQuantizer(
-                                                        **self._get_weights_configurable_quantizer_kwargs(n,
-                                                                                                          kernel_attr))})
+                wq = ConfigurableWeightsQuantizer(**self._get_weights_configurable_quantizer_kwargs(n, kernel_attr))
+                return KerasQuantizationWrapper(layer, weights_quantizers={kernel_attr: wq})
             else:
                 # TODO: Do we want to include other quantized attributes that are not
                 #  the kernel attribute in the mixed precision model?
@@ -106,12 +103,12 @@ class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
                 if not len(node_weights_qc) == 1:
                     Logger.critical(f"Expected a unique weights configuration for node {n.name}, but found {len(node_weights_qc)} configurations.")# pragma: no cover
 
+                weights_quant_cfg = node_weights_qc[0].weights_quantization_cfg
+                weights_quant_method = weights_quant_cfg.get_attr_config(kernel_attr).weights_quantization_method
                 quantier_for_node = get_inferable_quantizer_class(QuantizationTarget.Weights,
-                                                                  node_weights_qc[0].weights_quantization_cfg
-                                                                  .get_attr_config(kernel_attr)
-                                                                  .weights_quantization_method,
+                                                                  weights_quant_method,
                                                                   BaseKerasInferableQuantizer)
-                kwargs = get_inferable_quantizer_kwargs(node_weights_qc[0].weights_quantization_cfg,
+                kwargs = get_inferable_quantizer_kwargs(weights_quant_cfg,
                                                         QuantizationTarget.Weights,
                                                         kernel_attr)
 

{mct-nightly-2.2.0.20241202.131715 → mct-nightly-2.2.0.20241204.524}/model_compression_toolkit/core/keras/keras_implementation.py

@@ -438,17 +438,11 @@ class KerasImplementation(FrameworkImplementation):
             node: Node to indicate whether it needs to be part of the interest points set.
         Returns: True if the node should be considered an interest point, False otherwise.
         """
-
-        if node.is_match_type(Activation):
-            node_type_name = node.framework_attr[keras_constants.ACTIVATION]
-            if node_type_name in [keras_constants.SOFTMAX, keras_constants.SIGMOID]:
-                return True
-        elif any([node.is_match_type(_type) for _type in [tf.nn.softmax, tf.keras.layers.Softmax, tf.nn.sigmoid, Conv2D,
-                                                          DepthwiseConv2D, Conv2DTranspose, Dense, Concatenate, tf.concat,
-                                                          Add, tf.add]]):
+        if self.is_softmax(node) or self.is_sigmoid(node):
             return True
 
-        return False
+        return any([node.is_match_type(_type) for _type in [Conv2D, DepthwiseConv2D, Conv2DTranspose, Dense,
+                                                            Concatenate, tf.concat, Add, tf.add]])
 
     def get_mp_node_distance_fn(self, n: BaseNode,
                                 compute_distance_fn: Callable = None,
@@ -466,32 +460,34 @@ class KerasImplementation(FrameworkImplementation):
         Returns: A distance function between two tensors and a axis on which the distance is computed (if exists).
         """
 
-        axis = n.framework_attr.get(keras_constants.AXIS) \
-            if not isinstance(n, FunctionalNode) else n.op_call_kwargs.get(keras_constants.AXIS)
-
-        layer_class = n.layer_class
-        framework_attrs = n.framework_attr
+        axis = n.op_call_kwargs.get(keras_constants.AXIS) if isinstance(n, FunctionalNode) else n.framework_attr.get(keras_constants.AXIS)
 
         if compute_distance_fn is not None:
             return compute_distance_fn, axis
 
-        if layer_class == Activation:
-            node_type_name = framework_attrs[ACTIVATION]
-            if node_type_name == SOFTMAX and axis is not None:
-                return compute_kl_divergence, axis
-            elif node_type_name == SIGMOID:
-                return compute_cs, axis
-        elif axis is not None and (layer_class == tf.nn.softmax or layer_class == tf.keras.layers.Softmax
-                                   or (layer_class == TFOpLambda and
-                                       SOFTMAX in framework_attrs[keras_constants.FUNCTION])):
+        # TODO should we really return mse if axis is None? Error? Fill default?
+        if self.is_softmax(n) and axis is not None:
             return compute_kl_divergence, axis
-        elif layer_class == tf.nn.sigmoid or (layer_class == TFOpLambda and
-                                              SIGMOID in framework_attrs[keras_constants.FUNCTION]):
-            return compute_cs, axis
-        elif layer_class == Dense:
+
+        if self.is_sigmoid(n) or n.layer_class == Dense:
             return compute_cs, axis
+
         return partial(compute_mse, norm=norm_mse), axis
 
+    @staticmethod
+    def is_sigmoid(node: BaseNode):
+        cls = node.layer_class
+        return ((cls == Activation and node.framework_attr[ACTIVATION] == SIGMOID) or
+                cls == tf.nn.sigmoid or
+                cls == TFOpLambda and SIGMOID in node.framework_attr[keras_constants.FUNCTION])
+
+    @staticmethod
+    def is_softmax(node: BaseNode):
+        cls = node.layer_class
+        return ((cls == Activation and node.framework_attr[ACTIVATION] == SOFTMAX) or
+                cls in [tf.nn.softmax, tf.keras.layers.Softmax] or
+                cls == TFOpLambda and SOFTMAX in node.framework_attr[keras_constants.FUNCTION])
+
     def get_hessian_scores_calculator(self,
                                       graph: Graph,
                                       input_images: List[Any],

{mct-nightly-2.2.0.20241202.131715 → mct-nightly-2.2.0.20241204.524}/model_compression_toolkit/core/pytorch/pytorch_implementation.py

@@ -427,10 +427,8 @@ class PytorchImplementation(FrameworkImplementation):
         Returns: True if the node should be considered an interest point, False otherwise.
         """
 
-        if any([node.is_match_type(_type) for _type in [Conv2d, Linear, ConvTranspose2d, Sigmoid, sigmoid, Softmax,
-                                                        softmax, operator.add, add, cat, operator.concat]]):
-            return True
-        return False
+        return any(node.is_match_type(_type) for _type in [Conv2d, Linear, ConvTranspose2d, Sigmoid, sigmoid, Softmax,
+                                                           softmax, operator.add, add, cat, operator.concat])
 
     def get_mp_node_distance_fn(self, n: BaseNode,
                                 compute_distance_fn: Callable = None,

mct-nightly-2.2.0.20241202.131715/model_compression_toolkit/core/common/mixed_precision/mixed_precision_candidates_filter.py

@@ -1,80 +0,0 @@
-# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-import numpy as np
-
-from model_compression_toolkit.core import ResourceUtilization, FrameworkInfo
-from model_compression_toolkit.core.common import Graph
-from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
-
-
-def filter_candidates_for_mixed_precision(graph: Graph,
-                                          target_resource_utilization: ResourceUtilization,
-                                          fw_info: FrameworkInfo,
-                                          tpc: TargetPlatformCapabilities):
-    """
-    Filters out candidates in case of mixed precision search for only weights or activation compression.
-    For instance, if running only weights compression - filters out candidates of activation configurable nodes
-    such that only a single candidate would remain, with the bitwidth equal to the one defined in the matching layer's
-    base config in the TPC.
-
-    Note" This function modifies the graph inplace!
-
-    Args:
-        graph: A graph representation of the model to be quantized.
-        target_resource_utilization: The resource utilization of the target device.
-        fw_info: fw_info: Information needed for quantization about the specific framework.
-        tpc: TargetPlatformCapabilities object that describes the desired inference target platform.
-
-    """
-
-    no_total_restrictions = (target_resource_utilization.total_memory == np.inf and
-                             target_resource_utilization.bops == np.inf)
-
-    if target_resource_utilization.weights_memory < np.inf:
-        if target_resource_utilization.activation_memory == np.inf and no_total_restrictions:
-            # Running mixed precision for weights compression only -
-            # filter out candidates activation only configurable node
-            weights_conf = graph.get_weights_configurable_nodes(fw_info)
-            for n in graph.get_activation_configurable_nodes():
-                if n not in weights_conf:
-                    base_cfg_nbits = n.get_qco(tpc).base_config.activation_n_bits
-                    filtered_conf = [c for c in n.candidates_quantization_cfg if
-                                     c.activation_quantization_cfg.enable_activation_quantization and
-                                     c.activation_quantization_cfg.activation_n_bits == base_cfg_nbits]
-
-                    if len(filtered_conf) != 1:
-                        Logger.critical(f"Running weights only mixed precision failed on layer {n.name} with multiple "
-                                        f"activation quantization configurations.")  # pragma: no cover
-                    n.candidates_quantization_cfg = filtered_conf
-
-    elif target_resource_utilization.activation_memory < np.inf:
-        if target_resource_utilization.weights_memory == np.inf and no_total_restrictions:
-            # Running mixed precision for activation compression only -
-            # filter out candidates weights only configurable node
-            activation_conf = graph.get_activation_configurable_nodes()
-            for n in graph.get_weights_configurable_nodes(fw_info):
-                if n not in activation_conf:
-                    kernel_attr = graph.fw_info.get_kernel_op_attributes(n.type)[0]
-                    base_cfg_nbits = n.get_qco(tpc).base_config.attr_weights_configs_mapping[kernel_attr].weights_n_bits
-                    filtered_conf = [c for c in n.candidates_quantization_cfg if
-                                     c.weights_quantization_cfg.get_attr_config(
-                                         kernel_attr).enable_weights_quantization and
-                                     c.weights_quantization_cfg.get_attr_config(
-                                         kernel_attr).weights_n_bits == base_cfg_nbits]
-                    if len(filtered_conf) != 1:
-                        Logger.critical(f"Running activation only mixed precision failed on layer {n.name} with multiple "
-                                        f"weights quantization configurations.")  # pragma: no cover
-                    n.candidates_quantization_cfg = filtered_conf