mct-nightly 2.2.0.20241202.131715__py3-none-any.whl → 2.2.0.20241204.524__py3-none-any.whl

{mct_nightly-2.2.0.20241202.131715.dist-info → mct_nightly-2.2.0.20241204.524.dist-info}/METADATA

@@ -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.dist-info → mct_nightly-2.2.0.20241204.524.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=Y0S0cGZR8YyEXFmrpzm6m5TafU10MMyI_bHB3477yZE,1573
+model_compression_toolkit/__init__.py,sha256=CwUJhq62PjrcRJgkwwmu5dArRV2bz7lgnxc2ebnm840,1573
 model_compression_toolkit/constants.py,sha256=i_R6uXBfO1ph_X6DNJych2x59SUojfJbn7dNjs_mZnc,3846
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
@@ -32,8 +32,8 @@ model_compression_toolkit/core/common/fusion/__init__.py,sha256=Rf1RcYmelmdZmBV5
 model_compression_toolkit/core/common/fusion/graph_fuser.py,sha256=8seu9jBpC7HartP1nJd7S_SYFICyemVpDV9ZJ0QUQ7E,6212
 model_compression_toolkit/core/common/fusion/layer_fusing.py,sha256=lOubqpc18TslhXZijWUJQAa1c3jIB2S-M-5HK78wJPQ,5548
 model_compression_toolkit/core/common/graph/__init__.py,sha256=Xr-Lt_qXMdrCnnOaUS_OJP_3iTTGfPCLf8_vSrQgCs0,773
-model_compression_toolkit/core/common/graph/base_graph.py,sha256=lg5QaBkRbmvM3tGZ0Q34S3m0CbFql3LUv5BaXLe5TG8,37824
-model_compression_toolkit/core/common/graph/base_node.py,sha256=UygLaWvpd_-rXVn-QdPORHRaHCywaemokgbhy9MK52E,31837
+model_compression_toolkit/core/common/graph/base_graph.py,sha256=GG13PAtndsMjIqINfrCN6llVkFrg5CBfij4z99ntieU,37815
+model_compression_toolkit/core/common/graph/base_node.py,sha256=mGiDcHnL5KybDYSiONSWtjrHOXI6tjjfACfjv2Piogc,31756
 model_compression_toolkit/core/common/graph/edge.py,sha256=buoSEUZwilWBK3WeBKpJ-GeDaUA1SDdOHxDpxU_bGpk,3784
 model_compression_toolkit/core/common/graph/functional_node.py,sha256=GH5wStmw8SoAj5IdT_-ItN1Meo_P5NUTt_5bgJC4fak,3935
 model_compression_toolkit/core/common/graph/graph_matchers.py,sha256=CrDoHYq4iPaflgJWmoJ1K4ziLrRogJvFTVWg8P0UcDU,4744
@@ -63,21 +63,21 @@ model_compression_toolkit/core/common/mixed_precision/bit_width_setter.py,sha256
 model_compression_toolkit/core/common/mixed_precision/configurable_quant_id.py,sha256=LLDguK7afsbN742ucLpmJr5TUfTyFpK1vbf2bpVr1v0,882
 model_compression_toolkit/core/common/mixed_precision/configurable_quantizer_utils.py,sha256=7dKMi5S0zQZ16m8NWn1XIuoXsKuZUg64G4-uK8-j1PQ,5177
 model_compression_toolkit/core/common/mixed_precision/distance_weighting.py,sha256=-x8edUyudu1EAEM66AuXPtgayLpzbxoLNubfEbFM5kU,2867
-model_compression_toolkit/core/common/mixed_precision/mixed_precision_candidates_filter.py,sha256=klmaMQDeFc3IxRLf6YX4Dw1opFksbLyN10yFHdKAtLo,4875
+model_compression_toolkit/core/common/mixed_precision/mixed_precision_candidates_filter.py,sha256=nZb0_inMxgqlx5lG6uMF0sskHR-5zMSClzxYALBeqLA,4531
 model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py,sha256=r1t025_QHshyoop-PZvL7x6UuXaeplCCU3h4VNBhJHo,4309
-model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py,sha256=pk8HRoShDhiUprBC4m1AFQv1SacS4hOrj0MRdbq-5gY,7556
-model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py,sha256=TTTux4YiOnQqt-2h7Y38959XaDwNZc0eufLMx_yws5U,37578
-model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py,sha256=CYp2VuxXf95lYivolIuVRjAyaY5dFsDn2qh8ZhTmc9A,27525
+model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py,sha256=B7xLl8P5eCz0_fBxocDlNiv6k-3MdfMUk2GjYKl2p5k,7522
+model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py,sha256=hlaV4ybreT0DR4ftLSPg5KTit3BEm9dWA7Y8NHpEJ8w,37532
+model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py,sha256=adjuvrJcan7Ua3nYlJX7T6qGkCRHGqWMaM5-099a9Us,27220
 model_compression_toolkit/core/common/mixed_precision/set_layer_to_bitwidth.py,sha256=P8QtKgFXtt5b2RoubzI5OGlCfbEfZsAirjyrkFzK26A,2846
 model_compression_toolkit/core/common/mixed_precision/solution_refinement_procedure.py,sha256=KifDMbm7qkSfvSl6pcZzQ82naIXzeKL6aT-VsvWZYyc,7901
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
-model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization.py,sha256=HILF7CIn-GYPvPmTFyvjWLhuLDwSGwdBcAaKFgVYrwk,4745
+model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization.py,sha256=MtPkZfPIJWI191Hbjp6JluUyLnqiJRi3zNf-CqVNuag,5053
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_data.py,sha256=sFuUgWwC0aEBpf9zWmCTIcAbykBj3t5vmWAoB_BjYWA,14979
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/ru_aggregation_methods.py,sha256=ttc8wPa_9LZansutQ2f1ss-RTzgTv739wy3qsdLzyyk,4217
-model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/ru_functions_mapping.py,sha256=QhuqaECEGLnYC08iD6-2XXcU7NXbPzYf1sQcjYlGak8,1682
+model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/ru_functions_mapping.py,sha256=mOxZwOQYnOwSJMiapEEH9o-89ujJdPxSl8zXpnApc0U,1850
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/ru_methods.py,sha256=WC1EHoNuo_lrzy4NRhGJ1cgmJ2IsFsbmP86mrVO3AVA,21506
 model_compression_toolkit/core/common/mixed_precision/search_methods/__init__.py,sha256=sw7LOPN1bM82o3SkMaklyH0jw-TLGK0-fl2Wq73rffI,697
-model_compression_toolkit/core/common/mixed_precision/search_methods/linear_programming.py,sha256=7ceagUJVk3GgLrD8HnmXMgZTRLjEvZW_RnzDR3ahNDM,16592
+model_compression_toolkit/core/common/mixed_precision/search_methods/linear_programming.py,sha256=QZlQtvAUXUNNc6H2mKEFZhQ-fjP1QCIsxsS5BrhaXvU,16547
 model_compression_toolkit/core/common/network_editors/__init__.py,sha256=vZmu55bYqiaOQs3AjfwWDXHmuKZcLHt-wm7uR5fPEqg,1307
 model_compression_toolkit/core/common/network_editors/actions.py,sha256=nid0_j-Cn10xvmztT8yCKW_6uA7JEnom9SW9syx7wc0,19594
 model_compression_toolkit/core/common/network_editors/edit_network.py,sha256=dfgawi-nB0ocAJ0xcGn9E-Zv203oUnQLuMiXpX8vTgA,1748
@@ -157,7 +157,7 @@ model_compression_toolkit/core/keras/constants.py,sha256=dh4elQWt6Q6NYRht5k5RiiO
 model_compression_toolkit/core/keras/custom_layer_validation.py,sha256=f-b14wuiIgitBe7d0MmofYhDCTO3IhwJgwrh-Hq_t_U,1192
 model_compression_toolkit/core/keras/data_util.py,sha256=JdomIJZfep0QYPtx2jlg0xJ40cd9S_I7BakaWQi0wKw,2681
 model_compression_toolkit/core/keras/default_framework_info.py,sha256=PYcER89eEXjKtR0T7-2Y4f7cckqoD5OQbpHePoRkMec,5030
-model_compression_toolkit/core/keras/keras_implementation.py,sha256=Hi8seiFJdFqgYGGC003Y4879JQ7rmVZe8YiJ76T7FDE,32133
+model_compression_toolkit/core/keras/keras_implementation.py,sha256=HwbIR7x4t-TBNbWHVvVNFk8z-KFt6zM0LWAUXQuNZrk,31753
 model_compression_toolkit/core/keras/keras_model_validation.py,sha256=1wNV2clFdC9BzIELRLSO2uKf0xqjLqlkTJudwtCeaJk,1722
 model_compression_toolkit/core/keras/keras_node_prior_info.py,sha256=HUmzEXDQ8LGX7uOYSRiLZ2TNbYxLX9J9IeAa6QYlifg,3927
 model_compression_toolkit/core/keras/resource_utilization_data_facade.py,sha256=s56UIgiPipUQRNd2sd1xW6GFfYNMBmrocRCNtvpYLbY,4977
@@ -167,7 +167,7 @@ model_compression_toolkit/core/keras/back2framework/factory_model_builder.py,sha
 model_compression_toolkit/core/keras/back2framework/float_model_builder.py,sha256=9SFHhX-JnkB8PvYIIHRYlReBDI_RkZY9LditzW_ElLk,2444
 model_compression_toolkit/core/keras/back2framework/instance_builder.py,sha256=fBj13c6zkVoWX4JJG18_uXPptiEJqXClE_zFbaFB6Q8,4517
 model_compression_toolkit/core/keras/back2framework/keras_model_builder.py,sha256=ns3zFjngea7yZHrxqcV5FQCAfGcstho37D2S2KQZpwE,18444
-model_compression_toolkit/core/keras/back2framework/mixed_precision_model_builder.py,sha256=ygIS1WIiftF1VC3oGhc8N6j7MryKtWgEg8nr50p7f4U,15587
+model_compression_toolkit/core/keras/back2framework/mixed_precision_model_builder.py,sha256=emsaCYyZBF7oQfXAR0edU7idiMInXLXRuGPcrUp4slM,15301
 model_compression_toolkit/core/keras/back2framework/quantized_model_builder.py,sha256=5wFb4nx_F0Wu4c8pLf6n6OzxOHtpOJ6_3mQsNSXIudU,2481
 model_compression_toolkit/core/keras/graph_substitutions/__init__.py,sha256=mjbqLD-KcG3eNeCYpu1GBS7VclGVOQ63x2p6mAAuba4,698
 model_compression_toolkit/core/keras/graph_substitutions/substitutions/__init__.py,sha256=mjbqLD-KcG3eNeCYpu1GBS7VclGVOQ63x2p6mAAuba4,698
@@ -223,7 +223,7 @@ model_compression_toolkit/core/pytorch/constants.py,sha256=YwD_joIF0vK8UG2vW1NVv
 model_compression_toolkit/core/pytorch/data_util.py,sha256=YYbT135HhlTt0q6XdD2JX7AS_L92f_uV2rWq2hsJOCA,6325
 model_compression_toolkit/core/pytorch/default_framework_info.py,sha256=-Vls1P_8Ckm_18nnOsmQkZ71SmzHwtQLbQ383Z4Rb-U,4365
 model_compression_toolkit/core/pytorch/pytorch_device_config.py,sha256=S25cuw10AW3SEN_fRAGRcG_I3wdvvQx1ehSJzPnn-UI,4404
-model_compression_toolkit/core/pytorch/pytorch_implementation.py,sha256=SgxmSdzAQOPI9YHt4Q9-OeDi8fzAdgASHQ4nZ5maPsg,29599
+model_compression_toolkit/core/pytorch/pytorch_implementation.py,sha256=Xwt7eHS-QJJc1fyOrxL2tz8E2CP-b2M0_R-Dgb1Gm-4,29558
 model_compression_toolkit/core/pytorch/pytorch_node_prior_info.py,sha256=2LDQ7qupglHQ7o1Am7LWdfYVacfQnl-aW2N6l9det1w,3264
 model_compression_toolkit/core/pytorch/resource_utilization_data_facade.py,sha256=xpKj99OZKT9NT0vKIl_cOe8d89d2gef1gKoNT6PFElE,4989
 model_compression_toolkit/core/pytorch/utils.py,sha256=7VbgcLwtQvdEEc_AJgSOQ3U3KRKCICFPaBirN1fIQxg,3940
@@ -559,8 +559,8 @@ model_compression_toolkit/xquant/pytorch/model_analyzer.py,sha256=b93o800yVB3Z-i
 model_compression_toolkit/xquant/pytorch/pytorch_report_utils.py,sha256=bOc-hFL3gdoSM1Th_S2N_-9JJSlPGpZCTx_QLJHS6lg,3388
 model_compression_toolkit/xquant/pytorch/similarity_functions.py,sha256=CERxq5K8rqaiE-DlwhZBTUd9x69dtYJlkHOPLB54vm8,2354
 model_compression_toolkit/xquant/pytorch/tensorboard_utils.py,sha256=mkoEktLFFHtEKzzFRn_jCnxjhJolK12TZ5AQeDHzUO8,9767
-mct_nightly-2.2.0.20241202.131715.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.2.0.20241202.131715.dist-info/METADATA,sha256=YvbvKkI18wDQ20u0NKp5GtD5gZWfoWMYEXcYbMbNNeg,26449
-mct_nightly-2.2.0.20241202.131715.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-mct_nightly-2.2.0.20241202.131715.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.2.0.20241202.131715.dist-info/RECORD,,
+mct_nightly-2.2.0.20241204.524.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.2.0.20241204.524.dist-info/METADATA,sha256=O3ETKzNDjZGmSvp_WVmqIJz-jyk93WLG676QjyRsISs,26446
+mct_nightly-2.2.0.20241204.524.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+mct_nightly-2.2.0.20241204.524.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.2.0.20241204.524.dist-info/RECORD,,

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"

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]:
         """

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):
         """

model_compression_toolkit/core/common/mixed_precision/mixed_precision_candidates_filter.py

@@ -40,41 +40,38 @@ def filter_candidates_for_mixed_precision(graph: Graph,
 
     """
 
-    no_total_restrictions = (target_resource_utilization.total_memory == np.inf and
-                             target_resource_utilization.bops == np.inf)
+    tru = target_resource_utilization
+    if tru.total_mem_restricted() or tru.bops_restricted():
+        return
 
-    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 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
+            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
+    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

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)

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]

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.

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)}

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(

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)
 

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)
 

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],

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,