mct-nightly 2.3.0.20250512.625__py3-none-any.whl → 2.3.0.20250513.611__py3-none-any.whl

{mct_nightly-2.3.0.20250512.625.dist-info → mct_nightly-2.3.0.20250513.611.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mct-nightly
-Version: 2.3.0.20250512.625
+Version: 2.3.0.20250513.611
 Summary: A Model Compression Toolkit for neural networks
 Author-email: ssi-dnn-dev@sony.com
 Classifier: Programming Language :: Python :: 3

{mct_nightly-2.3.0.20250512.625.dist-info → mct_nightly-2.3.0.20250513.611.dist-info}/RECORD

@@ -1,5 +1,5 @@
-mct_nightly-2.3.0.20250512.625.dist-info/licenses/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-model_compression_toolkit/__init__.py,sha256=eq4dWyKngj1OrsAWFhmI9beCZaLAWyvU5DgVpuAMXK4,1557
+mct_nightly-2.3.0.20250513.611.dist-info/licenses/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+model_compression_toolkit/__init__.py,sha256=p_G6GkwHl_GiPtc0E2qL6iUBG-UpYcgFx1HDi073s0Q,1557
 model_compression_toolkit/constants.py,sha256=iJ6vfTjC2oFIZWt8wvHoxEw5YJi3yl0Hd4q30_8q0Zc,3958
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=L3q7tn3Uht0i_7phnlOWMR2Te2zvzrt2HOz9vYEInts,4529
@@ -12,14 +12,14 @@ model_compression_toolkit/core/quantization_prep_runner.py,sha256=DPevqQ8brkdut8
 model_compression_toolkit/core/runner.py,sha256=_r6cieb7Ur2BeHQK5XxTZHogjyA0utybvIVbH06CBHY,13056
 model_compression_toolkit/core/common/__init__.py,sha256=Wh127PbXcETZX_d1PQqZ71ETK3J9XO5A-HpadGUbj6o,1447
 model_compression_toolkit/core/common/base_substitutions.py,sha256=xDFSmVVs_iFSZfajytI0cuQaNRNcwHX3uqOoHgVUvxQ,1666
-model_compression_toolkit/core/common/framework_implementation.py,sha256=L88uv_sfYM_56FSmxXP--emjv01_lk7IPqOI7QBZEt0,22939
+model_compression_toolkit/core/common/framework_implementation.py,sha256=JQI_eoZZoNk5Y_jAxLfYt9-wzfs7zGpTldz9UblxmMc,21182
 model_compression_toolkit/core/common/framework_info.py,sha256=5tderHT-7Cd21QrRFIJj3hH_gAcnlivOzwZ5m1ldJOs,6526
 model_compression_toolkit/core/common/memory_computation.py,sha256=ixoSpV5ZYZGyzhre3kQcvR2sNA8KBsPZ3lgbkDnw9Cs,1205
 model_compression_toolkit/core/common/model_builder_mode.py,sha256=jll9-59OPaE3ug7Y9-lLyV99_FoNHxkGZMgcm0Vkpss,1324
 model_compression_toolkit/core/common/model_collector.py,sha256=Tno3-qx9jmPZAZyLYgbPlMLHakVfuEH5deuToZNuCb0,13195
 model_compression_toolkit/core/common/model_validation.py,sha256=LaG8wd6aZl0OJgieE3SeiVDEPxtk8IHq9-3wSnmWhY4,1214
 model_compression_toolkit/core/common/node_prior_info.py,sha256=WXX_PrGVG9M9I_REG5ZzFBohwmV4yf356sZnrja_FLo,2832
-model_compression_toolkit/core/common/similarity_analyzer.py,sha256=FikcIqgQQpfiXr9VJvgl-wk8OyH7-LvC8ku7TkhJfJM,9200
+model_compression_toolkit/core/common/similarity_analyzer.py,sha256=S3f6WgHyw62dGcxpX51FGKyfebe2zv9ABKbjtGyKRvY,9215
 model_compression_toolkit/core/common/user_info.py,sha256=dSRMnT-oewmdOziIpEuW-s9K7vTSeyUBxT4z9neXurI,1648
 model_compression_toolkit/core/common/back2framework/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/core/common/back2framework/base_model_builder.py,sha256=V1oShKzbSkdcTvREn8VnQQBzvm-tTHkWMXqMkYozF2s,2023
@@ -66,13 +66,13 @@ model_compression_toolkit/core/common/mixed_precision/configurable_quant_id.py,s
 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=6pLUEEIqRTVIlCYQC4JIvY55KAvuBHEX8uTOQ-1Ac4Q,3859
-model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py,sha256=onHgDwfw8CUbZFNU-RYit9eqA6FrzAtFA3akVZ2d7IM,4533
+model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py,sha256=BO4ouM_UVS9Fg0z95gLJSMz1ep6YQC5za_iXI_qW2yQ,5399
 model_compression_toolkit/core/common/mixed_precision/mixed_precision_ru_helper.py,sha256=-hOMBucYn12ePyLd0b1KxniPOIRu4b53SwEzv0bWToI,4943
-model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py,sha256=d5-3j2e_rdcQOT7c4s0p7640i3nSetjJ6MgMhhMM7dc,6152
-model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py,sha256=Lk5cftihGpgFQoyqnRGiwJFFqkI8dkx0l1q0sVJi2CE,27505
-model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py,sha256=R3UIO9lKf-lpEGfJOqgpQAXdP1IWMatWxXKYDkhWj_E,28096
-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=S1ChgxtUjzXJufNWyRbKoNdyNC6fGUjPeComDMx8ZCo,9479
+model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py,sha256=KhiHGpmN5QbpyJQnTZmXigdXFlSlRNqpOOyKGj1Fwek,6412
+model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py,sha256=OzRhlJ2IS9Dwv0rgobee0xTtAeRwlBC6KvVEcx2_oB0,28089
+model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py,sha256=4uhUXKgwyMrJqEVK5uJzVr67GI5YzDTHLveV4maB7z0,28079
+model_compression_toolkit/core/common/mixed_precision/set_layer_to_bitwidth.py,sha256=Zn6SgzGLWWKmuYGHd1YtKxZdYnQWRDeXEkKlBiTbHcs,2929
+model_compression_toolkit/core/common/mixed_precision/solution_refinement_procedure.py,sha256=MY8df-c_kITEr_7hOctaxhdiq29hSTA0La9Qo0oTJJY,9678
 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=PKkhc5q8pEPnNLXwo3U56EOCfYnPXIvPs0LlCGZOoKU,4426
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_calculator.py,sha256=-kNcmQQFVHRPizInaRrCEIuh_q_57CWxC6CIV6azF4g,39640
@@ -106,7 +106,7 @@ model_compression_toolkit/core/common/quantization/candidate_node_quantization_c
 model_compression_toolkit/core/common/quantization/core_config.py,sha256=yxCzWqldcHoe8GGxrH0tp99bhrc5jDT7SgZftnMUUBE,2374
 model_compression_toolkit/core/common/quantization/debug_config.py,sha256=uH45Uq3Tp9FIyMynex_WY2_y-Kv8LuPw2XXZydnpW5A,1649
 model_compression_toolkit/core/common/quantization/filter_nodes_candidates.py,sha256=n2A8pO7_DMMae4o69U0I00iW6mzeRlRfKHDxlQUBBuI,7204
-model_compression_toolkit/core/common/quantization/node_quantization_config.py,sha256=0OJZtQuv-StbKZOpalvGi9lcpHJNRPeuclevSaCPggc,29792
+model_compression_toolkit/core/common/quantization/node_quantization_config.py,sha256=XmuG1ZwBxZEEcVKBwCo_v3vfjVjIJqtyO94QCeczddw,30131
 model_compression_toolkit/core/common/quantization/quantization_config.py,sha256=UkSVW7d1OF_Px9gAjsqqK65aYhIBFWaBO-_IH6_AFfg,4403
 model_compression_toolkit/core/common/quantization/quantization_fn_selection.py,sha256=HfBkSiRTOf9mNF-TNQHTCCs3xSg66F20no0O6vl5v1Y,2154
 model_compression_toolkit/core/common/quantization/quantization_params_fn_selection.py,sha256=7eG7dl1TcbdnHwgmvyjarxLs0o6Lw_9VAjXAm4rsiBk,3791
@@ -158,7 +158,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=jm54o-SlI1DJ-sEvRuX9OyLN68tEt0VxcqrdIjR98Ag,8366
 model_compression_toolkit/core/keras/default_framework_info.py,sha256=DvK1Tr6z3cQlJw1nx62iFaeSsQSXJl55xOIcJ1uNGu8,5020
-model_compression_toolkit/core/keras/keras_implementation.py,sha256=_15BrSGTRSSp_8ayuo2x-hdKanew1xuIPSumP46IGSA,32545
+model_compression_toolkit/core/keras/keras_implementation.py,sha256=bQFnX6OjUBo3q4aEPrxbXigqSoIqAsm5YLoXWDvqghE,30048
 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=XBCmUrHy_fNQCfSjnXCpwuEtc7cda4hXySuiIzhFGqc,5696
@@ -168,7 +168,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=TY86-Mb8hmo8RgCcQvkSYIthYOqV9e4VIMpqIyouJ4Y,17397
-model_compression_toolkit/core/keras/back2framework/mixed_precision_model_builder.py,sha256=BTDJB6VUAyVapzkwnftdXkv9RaQfwp_GIEk1FyovdGg,14813
+model_compression_toolkit/core/keras/back2framework/mixed_precision_model_builder.py,sha256=s3uxI5X7q82L1Bui02Z16yeGjKBiFXjjJoPoBhUGO0Y,11804
 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
@@ -198,8 +198,8 @@ model_compression_toolkit/core/keras/hessian/activation_hessian_scores_calculato
 model_compression_toolkit/core/keras/hessian/hessian_scores_calculator_keras.py,sha256=1o7X9GXSfpEmuB5ee2AaBQ2sN2xzX4-smbrq_0qOGRU,4454
 model_compression_toolkit/core/keras/hessian/weights_hessian_scores_calculator_keras.py,sha256=Rl6NNGkHMV0ioEM5bbM4XX7yHDqG6mMp4ifN2VQBDxE,12168
 model_compression_toolkit/core/keras/mixed_precision/__init__.py,sha256=sw7LOPN1bM82o3SkMaklyH0jw-TLGK0-fl2Wq73rffI,697
-model_compression_toolkit/core/keras/mixed_precision/configurable_activation_quantizer.py,sha256=WFwPtCcXR3qY86OML_jyzasvdd2DGhy4-GveAGpDOt0,5075
-model_compression_toolkit/core/keras/mixed_precision/configurable_weights_quantizer.py,sha256=38Lvwux9L35oT6muck6_FH7nDdH2N8_kuGDMj4-QNpE,6647
+model_compression_toolkit/core/keras/mixed_precision/configurable_activation_quantizer.py,sha256=aB3y8vaiXSmy1bpvlqXDswL3-FTz019s5r-lcb4FKhE,5254
+model_compression_toolkit/core/keras/mixed_precision/configurable_weights_quantizer.py,sha256=GtW0yars8PzqP9uL_vfXrtqHwKiStmOxPng20rYaIjU,6805
 model_compression_toolkit/core/keras/pruning/__init__.py,sha256=3Lkr37Exk9u8811hw8hVqkGcbTQGcLjd3LLuLC3fa_E,698
 model_compression_toolkit/core/keras/pruning/pruning_keras_implementation.py,sha256=EJkblZ4OAjI5l29GKsUraam5Jn58Sogld47_rFFyr3k,12777
 model_compression_toolkit/core/keras/quantizer/__init__.py,sha256=mjbqLD-KcG3eNeCYpu1GBS7VclGVOQ63x2p6mAAuba4,698
@@ -224,7 +224,7 @@ model_compression_toolkit/core/pytorch/constants.py,sha256=Sg0hkUaMe88mI2_pd3Kqh
 model_compression_toolkit/core/pytorch/data_util.py,sha256=YYbT135HhlTt0q6XdD2JX7AS_L92f_uV2rWq2hsJOCA,6325
 model_compression_toolkit/core/pytorch/default_framework_info.py,sha256=-byHTXmQEuOiqTAX45BHGi3mRRBF4_EfJ3XhpmVilSU,4355
 model_compression_toolkit/core/pytorch/pytorch_device_config.py,sha256=S25cuw10AW3SEN_fRAGRcG_I3wdvvQx1ehSJzPnn-UI,4404
-model_compression_toolkit/core/pytorch/pytorch_implementation.py,sha256=c_QFo4e7t6b21CDakGhjVpqy5aXFxxqkdJ-s54HEOfs,31207
+model_compression_toolkit/core/pytorch/pytorch_implementation.py,sha256=33P9kttfgVKtbqOgmstSLuzs6MSwXaK8MbvtAScqiBI,28755
 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=aIHl-dTAC4ISnWSKLD99c-1W3827vfRGyLjMBib-l3s,5618
 model_compression_toolkit/core/pytorch/utils.py,sha256=xNVE7YMtHupLEimIJcxmfcMGM4XKB9I1v0-K8lDeLB8,3936
@@ -232,7 +232,7 @@ model_compression_toolkit/core/pytorch/back2framework/__init__.py,sha256=H_WixgN
 model_compression_toolkit/core/pytorch/back2framework/factory_model_builder.py,sha256=bwppTPRs6gL96nm7qPiKrNcBj4Krr0yEsOWjRF0aXmQ,2339
 model_compression_toolkit/core/pytorch/back2framework/float_model_builder.py,sha256=tLrlUyYhxVKVjkad1ZAtbRra0HedB3iVfIkZ_dYnQ-4,3419
 model_compression_toolkit/core/pytorch/back2framework/instance_builder.py,sha256=BBHBfTqeWm7L3iDyPBpk0jxvj-rBg1QWI23imkjfIl0,1467
-model_compression_toolkit/core/pytorch/back2framework/mixed_precision_model_builder.py,sha256=Pyk2hrwCB-9HyhZUqtKFK_t6YBL5yKB-S0CStfEhx_M,14675
+model_compression_toolkit/core/pytorch/back2framework/mixed_precision_model_builder.py,sha256=HAzzWOnPcIeDxQO1712254RNTBZD-gVSMSVnxqpfuQ0,11907
 model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py,sha256=Z-ZQV-GWdOBGPbksiWBQ8MtFkQ41qgUKU5d5c8aNSjQ,21646
 model_compression_toolkit/core/pytorch/back2framework/quantized_model_builder.py,sha256=qZNNOlNTTV4ZKPG3q5GDXkIVTPUEr8dvxAS_YiMORmg,3456
 model_compression_toolkit/core/pytorch/back2framework/quantization_wrapper/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
@@ -268,8 +268,8 @@ model_compression_toolkit/core/pytorch/hessian/activation_hessian_scores_calcula
 model_compression_toolkit/core/pytorch/hessian/hessian_scores_calculator_pytorch.py,sha256=8f_XlM8ZFVQPNGr1iECr1hv8QusYDrNU_vTkLQZE9RU,2477
 model_compression_toolkit/core/pytorch/hessian/weights_hessian_scores_calculator_pytorch.py,sha256=UzWxWDbr8koKZatEcPn8RCb0Zjm_7fKTvIGb98sp18k,8487
 model_compression_toolkit/core/pytorch/mixed_precision/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
-model_compression_toolkit/core/pytorch/mixed_precision/configurable_activation_quantizer.py,sha256=mcY_KOQgABIqGIMh0x6mNxaKp7SFNbkEIYavR2X7SQ4,4754
-model_compression_toolkit/core/pytorch/mixed_precision/configurable_weights_quantizer.py,sha256=zp1Xp75IDf9LN5YGO2UzeDbms_6ICQ_pSE1ORQr-SA8,6281
+model_compression_toolkit/core/pytorch/mixed_precision/configurable_activation_quantizer.py,sha256=DVQrEJbB7MKj_LitU92cBxDApwnAAkilYvQzkr79ffg,4813
+model_compression_toolkit/core/pytorch/mixed_precision/configurable_weights_quantizer.py,sha256=KDnwmbhvhJMfNg1IuTvvzBNEriPQH9bL9dJ5VvWTzpE,6631
 model_compression_toolkit/core/pytorch/pruning/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
 model_compression_toolkit/core/pytorch/pruning/pruning_pytorch_implementation.py,sha256=VfEEVwWEXKpVlZFnr7N6mvEjcpq85ROLg05ZvXfD1Pg,14764
 model_compression_toolkit/core/pytorch/quantizer/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
@@ -528,7 +528,7 @@ model_compression_toolkit/xquant/pytorch/model_analyzer.py,sha256=b93o800yVB3Z-i
 model_compression_toolkit/xquant/pytorch/pytorch_report_utils.py,sha256=UVN_S9ULHBEldBpShCOt8-soT8YTQ5oE362y96qF_FA,3950
 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.3.0.20250512.625.dist-info/METADATA,sha256=Rl3DUdbepLzTi6pJhMjSt65rk2rnMQgUiJYzRi-dCPk,25136
-mct_nightly-2.3.0.20250512.625.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
-mct_nightly-2.3.0.20250512.625.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.3.0.20250512.625.dist-info/RECORD,,
+mct_nightly-2.3.0.20250513.611.dist-info/METADATA,sha256=dx0fsYTzsB_Y1IVuSNMaJPgPO4lhotb3TlDZ-dq2JF8,25136
+mct_nightly-2.3.0.20250513.611.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+mct_nightly-2.3.0.20250513.611.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.3.0.20250513.611.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.3.0.20250512.000625"
+__version__ = "2.3.0.20250513.000611"

model_compression_toolkit/core/common/framework_implementation.py

@@ -13,24 +13,20 @@
 # limitations under the License.
 # ==============================================================================
 from abc import ABC, abstractmethod
-from typing import Callable, Any, List, Tuple, Dict, Generator
+from typing import Callable, Any, List, Tuple, Generator, Type
 
 import numpy as np
 
 from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS
-from model_compression_toolkit.core import MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode
-from model_compression_toolkit.core.common.collectors.statistics_collector import BaseStatsCollector
 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.hessian import HessianScoresRequest, HessianInfoService
-from model_compression_toolkit.core.common.mixed_precision.sensitivity_evaluation import SensitivityEvaluation
+from model_compression_toolkit.core.common.hessian import HessianScoresRequest
 from model_compression_toolkit.core.common.model_builder_mode import ModelBuilderMode
 from model_compression_toolkit.core.common.node_prior_info import NodePriorInfo
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig
-from model_compression_toolkit.core.common.user_info import UserInformation
 
 
 class FrameworkImplementation(ABC):
@@ -38,6 +34,10 @@ class FrameworkImplementation(ABC):
     An abstract class with abstract methods that should be implemented when supporting a new
     framework in MCT.
     """
+    weights_quant_layer_cls: Type
+    activation_quant_layer_cls: Type
+    configurable_weights_quantizer_cls: Type
+    configurable_activation_quantizer_cls: Type
 
     @property
     def constants(self):
@@ -327,33 +327,6 @@ class FrameworkImplementation(ABC):
                              f'framework\'s get_substitutions_after_second_moment_correction '
                              f'method.')  # pragma: no cover
 
-    @abstractmethod
-    def get_sensitivity_evaluator(self,
-                                  graph: Graph,
-                                  quant_config: MixedPrecisionQuantizationConfig,
-                                  representative_data_gen: Callable,
-                                  fw_info: FrameworkInfo,
-                                  hessian_info_service: HessianInfoService = None,
-                                  disable_activation_for_metric: bool = False) -> SensitivityEvaluation:
-        """
-        Creates and returns an object which handles the computation of a sensitivity metric for a mixed-precision
-        configuration (comparing to the float model).
-
-        Args:
-            graph: Graph to build its float and mixed-precision models.
-            quant_config: QuantizationConfig of how the model should be quantized.
-            representative_data_gen: Dataset to use for retrieving images for the models inputs.
-            fw_info: FrameworkInfo object with information about the specific framework's model.
-            disable_activation_for_metric: Whether to disable activation quantization when computing the MP metric.
-            hessian_info_service: HessianInfoService to fetch information based on Hessian-approximation.
-
-        Returns:
-            A function that computes the metric.
-        """
-
-        raise NotImplementedError(f'{self.__class__.__name__} has to implement the '
-                             f'framework\'s get_sensitivity_evaluator method.')  # pragma: no cover
-
     def get_node_prior_info(self, node: BaseNode,
                             fw_info: FrameworkInfo,
                             graph: Graph) -> NodePriorInfo:

model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py

@@ -14,11 +14,23 @@
 # ==============================================================================
 
 from dataclasses import dataclass, field
+from enum import Enum
 from typing import List, Callable, Optional
 from model_compression_toolkit.constants import MP_DEFAULT_NUM_SAMPLES, ACT_HESSIAN_DEFAULT_BATCH_SIZE
 from model_compression_toolkit.core.common.mixed_precision.distance_weighting import MpDistanceWeighting
 
 
+class MpMetricNormalization(Enum):
+    """
+    MAXBIT: normalize sensitivity metrics of layer candidates by max-bitwidth candidate (of that layer).
+    MINBIT: normalize sensitivity metrics of layer candidates by min-bitwidth candidate (of that layer).
+    NONE: no normalization.
+    """
+    MAXBIT = 'MAXBIT'
+    MINBIT = 'MINBIT'
+    NONE = 'NONE'
+
+
 @dataclass
 class MixedPrecisionQuantizationConfig:
     """
@@ -27,7 +39,6 @@ class MixedPrecisionQuantizationConfig:
     Args:
         compute_distance_fn (Callable): Function to compute a distance between two tensors. If None, using pre-defined distance methods based on the layer type for each layer.
         distance_weighting_method (MpDistanceWeighting): MpDistanceWeighting enum value that provides a function to use when weighting the distances among different layers when computing the sensitivity metric.
-        custom_metric_fn (Callable): Function to compute a custom metric. As input gets the model_mp and returns a float value for metric. If None, uses interest point metric.
         num_of_images (int): Number of images to use to evaluate the sensitivity of a mixed-precision model comparing to the float model.
         configuration_overwrite (List[int]): A list of integers that enables overwrite of mixed precision with a predefined one.
         num_interest_points_factor (float): A multiplication factor between zero and one (represents percentage) to reduce the number of interest points used to calculate the distance metric.
@@ -36,11 +47,16 @@ class MixedPrecisionQuantizationConfig:
         refine_mp_solution (bool): Whether to try to improve the final mixed-precision configuration using a greedy algorithm that searches layers to increase their bit-width, or not.
         metric_normalization_threshold (float): A threshold for checking the mixed precision distance metric values, In case of values larger than this threshold, the metric will be scaled to prevent numerical issues.
         hessian_batch_size (int): The Hessian computation batch size. used only if using mixed precision with Hessian-based objective.
-    """
+        metric_normalization (MpMetricNormalization): Metric normalization method.
+        metric_epsilon (float | None): ensure minimal distance between the metric for any non-max-bidwidth candidate
+          and a max-bitwidth candidate, i.e. metric(non-max-bitwidth) >= metric(max-bitwidth) + epsilon.
+          If none, the computed metrics are used as is.
+        custom_metric_fn (Callable): Function to compute a custom metric. As input gets the model_mp and returns a
+          float value for metric. If None, uses interest point metric.
 
+    """
     compute_distance_fn: Optional[Callable] = None
     distance_weighting_method: MpDistanceWeighting = MpDistanceWeighting.AVG
-    custom_metric_fn: Optional[Callable] = None
     num_of_images: int = MP_DEFAULT_NUM_SAMPLES
     configuration_overwrite: Optional[List[int]] = None
     num_interest_points_factor: float = field(default=1.0, metadata={"description": "Should be between 0.0 and 1.0"})
@@ -49,6 +65,9 @@ class MixedPrecisionQuantizationConfig:
     refine_mp_solution: bool = True
     metric_normalization_threshold: float = 1e10
     hessian_batch_size: int = ACT_HESSIAN_DEFAULT_BATCH_SIZE
+    metric_normalization: MpMetricNormalization = MpMetricNormalization.NONE
+    metric_epsilon: Optional[float] = 1e-6
+    custom_metric_fn: Optional[Callable] = None
     _is_mixed_precision_enabled: bool = field(init=False, default=False)
 
     def __post_init__(self):

model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py

@@ -25,6 +25,7 @@ from model_compression_toolkit.core.common.mixed_precision.mixed_precision_searc
     MixedPrecisionSearchManager
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import \
     ResourceUtilization
+from model_compression_toolkit.core.common.mixed_precision.sensitivity_evaluation import SensitivityEvaluation
 from model_compression_toolkit.core.common.mixed_precision.solution_refinement_procedure import \
     greedy_solution_refinement_procedure
 
@@ -78,11 +79,12 @@ def search_bit_width(graph: Graph,
 
     # 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)
-    se = fw_impl.get_sensitivity_evaluator(
+    se = SensitivityEvaluation(
         graph,
         mp_config,
         representative_data_gen=representative_data_gen,
         fw_info=fw_info,
+        fw_impl=fw_impl,
         disable_activation_for_metric=disable_activation_for_metric,
         hessian_info_service=hessian_info_service)
 
@@ -96,10 +98,11 @@ def search_bit_width(graph: Graph,
 
     # Search manager and LP are highly coupled, so LP search method was moved inside search manager.
     search_manager = MixedPrecisionSearchManager(graph,
-                                                 fw_info,
-                                                 fw_impl,
-                                                 se,
-                                                 target_resource_utilization)
+                                                 fw_info=fw_info,
+                                                 fw_impl=fw_impl,
+                                                 sensitivity_evaluator=se,
+                                                 target_resource_utilization=target_resource_utilization,
+                                                 mp_config=mp_config)
     nodes_bit_cfg = search_manager.search()
 
     graph.skip_validation_check = False

model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import os
+
 import itertools
 
 import copy
@@ -19,7 +21,7 @@ from collections import defaultdict
 
 from tqdm import tqdm
 
-from typing import Dict, List, Tuple, Optional
+from typing import Dict, List, Tuple, Optional, Set
 
 import numpy as np
 
@@ -40,6 +42,8 @@ from model_compression_toolkit.core.common.mixed_precision.search_methods.linear
 from model_compression_toolkit.core.common.mixed_precision.sensitivity_evaluation import SensitivityEvaluation
 from model_compression_toolkit.core.common.substitutions.apply_substitutions import substitute
 from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
+    MixedPrecisionQuantizationConfig, MpMetricNormalization
 
 
 class MixedPrecisionSearchManager:
@@ -52,7 +56,8 @@ class MixedPrecisionSearchManager:
                  fw_info: FrameworkInfo,
                  fw_impl: FrameworkImplementation,
                  sensitivity_evaluator: SensitivityEvaluation,
-                 target_resource_utilization: ResourceUtilization):
+                 target_resource_utilization: ResourceUtilization,
+                 mp_config: MixedPrecisionQuantizationConfig):
         """
 
         Args:
@@ -74,21 +79,21 @@ class MixedPrecisionSearchManager:
 
         self.sensitivity_evaluator = sensitivity_evaluator
         self.target_resource_utilization = target_resource_utilization
+        self.mp_config = mp_config
 
         self.mp_topo_configurable_nodes = self.mp_graph.get_configurable_sorted_nodes(fw_info)
 
         self.ru_targets = target_resource_utilization.get_restricted_targets()
-        self.ru_helper = MixedPrecisionRUHelper(self.original_graph, fw_info, fw_impl)
+        self.orig_graph_ru_helper = MixedPrecisionRUHelper(self.original_graph, fw_info, fw_impl)
 
         self.min_ru_config: Dict[BaseNode, int] = self.mp_graph.get_min_candidates_config(fw_info)
-        self.max_ru_config: Dict[BaseNode, int] = self.mp_graph.get_max_candidates_config(fw_info)
 
-        self.config_reconstruction_helper = ConfigReconstructionHelper(self.original_graph)
+        self.config_reconstructor = None
+        orig_min_config = self.min_ru_config
         if self.using_virtual_graph:
-            real_min_ru_config = self.config_reconstruction_helper.reconstruct_full_configuration(self.min_ru_config)
-            self.min_ru = self.ru_helper.compute_utilization(self.ru_targets, real_min_ru_config)
-        else:
-            self.min_ru = self.ru_helper.compute_utilization(self.ru_targets, self.min_ru_config)
+            self.config_reconstructor = ConfigReconstructionHelper(self.original_graph)
+            orig_min_config = self.config_reconstructor.reconstruct_full_configuration(self.min_ru_config)
+        self.min_ru = self.orig_graph_ru_helper.compute_utilization(self.ru_targets, orig_min_config)
 
     def search(self) -> Dict[BaseNode, int]:
         """
@@ -100,7 +105,7 @@ class MixedPrecisionSearchManager:
         mp_config = self._prepare_and_run_solver()
 
         if self.using_virtual_graph:
-            mp_config = self.config_reconstruction_helper.reconstruct_full_configuration(mp_config)
+            mp_config = self.config_reconstructor.reconstruct_full_configuration(mp_config)
 
         return mp_config
 
@@ -143,61 +148,64 @@ class MixedPrecisionSearchManager:
                              f"following targets: {unsatisfiable_targets}")
         return rel_target_ru
 
-    def _build_sensitivity_mapping(self, eps: float = 1e-6) -> Dict[BaseNode, List[float]]:
+    def _build_sensitivity_mapping(self) -> Dict[BaseNode, List[float]]:
         """
         This function measures the sensitivity of a change in a bitwidth of a layer on the entire model.
 
-        Args:
-            eps: if sensitivity for a non-max candidate is lower than for a max candidate, we set it to
-              sensitivity of a max candidate + epsilon.
-
         Returns:
             Mapping from nodes to their bitwidth candidates sensitivity.
         """
-
         Logger.info('Starting to evaluate metrics')
-
-        orig_sorted_nodes = self.original_graph.get_configurable_sorted_nodes(self.fw_info)
-
-        def topo_cfg(cfg: dict) -> list:
-            topo_cfg = [cfg[n] for n in orig_sorted_nodes]
-            assert len(topo_cfg) == len(cfg)
-            return topo_cfg
-
-        def compute_metric(cfg, node_idx=None, baseline_cfg=None):
-            return self.sensitivity_evaluator.compute_metric(topo_cfg(cfg),
-                                                             node_idx,
-                                                             topo_cfg(baseline_cfg) if baseline_cfg else None)
-
-        if self.using_virtual_graph:
-            origin_max_config = self.config_reconstruction_helper.reconstruct_full_configuration(self.max_ru_config)
-            max_config_value = compute_metric(origin_max_config)
-        else:
-            max_config_value = compute_metric(self.max_ru_config)
+        norm_method = self.mp_config.metric_normalization
+        eps = self.mp_config.metric_epsilon
+
+        verbose = 'VERBOSE_MP_METRIC' in os.environ
+
+        def normalize(node_candidates_metrics, max_ind):
+            if norm_method == MpMetricNormalization.NONE:
+                return node_candidates_metrics
+            if norm_method == MpMetricNormalization.MAXBIT:
+                ref_ind = max_ind
+            elif norm_method == MpMetricNormalization.MINBIT:
+                ref_ind = node.find_min_candidate_index()
+            else:  # pragma: no cover
+                raise ValueError(f'Unexpected MpMetricNormalization mode {norm_method}')
+            normalized_metrics = node_candidates_metrics / node_candidates_metrics[ref_ind]
+            if verbose and not np.array_equal(normalized_metrics, node_candidates_metrics):
+                print(f'{"normalized metric:":25}', candidates_sensitivity)
+            return normalized_metrics
+
+        def ensure_maxbit_minimal_metric(node_candidates_metrics, max_ind):
+            if eps is None:
+                return node_candidates_metrics
+            # We want maxbit configuration to have the minimal distance metric (so that optimization objective
+            # doesn't prefer lower bits). If we got a smaller metric for non-maxbit, we update it to metric(maxbit)+eps.
+            max_val = node_candidates_metrics[max_ind]
+            metrics = np.maximum(node_candidates_metrics, max_val + eps)
+            metrics[max_ind] = max_val
+            if verbose and not np.array_equal(metrics, node_candidates_metrics):
+                print(f'{"eps-adjusted metric:":25}', candidates_sensitivity)
+            return metrics
 
         layer_to_metrics_mapping = defaultdict(list)
         for node_idx, node in tqdm(enumerate(self.mp_topo_configurable_nodes)):
+            candidates_sensitivity = np.empty(len(node.candidates_quantization_cfg))
             for bitwidth_idx, _ in enumerate(node.candidates_quantization_cfg):
-                if self.max_ru_config[node] == bitwidth_idx:
-                    # This is a computation of the metric for the max configuration, assign pre-calculated value
-                    layer_to_metrics_mapping[node].append(max_config_value)
-                    continue
-
-                # Create a configuration that differs at one layer only from the baseline model
-                mp_model_configuration = self.max_ru_config.copy()
-                mp_model_configuration[node] = bitwidth_idx
-
-                # Build a distance matrix using the function we got from the framework implementation.
                 if self.using_virtual_graph:
-                    # Reconstructing original graph's configuration from virtual graph's configuration
-                    orig_mp_config = self.config_reconstruction_helper.reconstruct_full_configuration(mp_model_configuration)
-                    changed_nodes = [orig_sorted_nodes.index(n) for n, ind in orig_mp_config.items()
-                                     if origin_max_config[n] != ind]
-                    metric_value = compute_metric(orig_mp_config, changed_nodes, origin_max_config)
+                    a_cfg, w_cfg = self.config_reconstructor.reconstruct_separate_aw_configs({node: bitwidth_idx})
                 else:
-                    metric_value = compute_metric(mp_model_configuration, [node_idx], self.max_ru_config)
-                metric_value = max(metric_value, max_config_value + eps)
-                layer_to_metrics_mapping[node].append(metric_value)
+                    a_cfg = {node: bitwidth_idx} if node.has_configurable_activation() else {}
+                    w_cfg = {node: bitwidth_idx} if node.has_any_configurable_weight() else {}
+                candidates_sensitivity[bitwidth_idx] = self.sensitivity_evaluator.compute_metric(
+                    mp_a_cfg={n.name: ind for n, ind in a_cfg.items()},
+                    mp_w_cfg={n.name: ind for n, ind in w_cfg.items()}
+                )
+            if verbose:
+                print(f'{node.name}\n{"raw metric:":25}', candidates_sensitivity)
+            max_ind = node.find_max_candidate_index()
+            candidates_sensitivity = normalize(candidates_sensitivity, max_ind)
+            candidates_sensitivity = ensure_maxbit_minimal_metric(candidates_sensitivity, max_ind)
+            layer_to_metrics_mapping[node] = candidates_sensitivity
 
         # Finalize distance metric mapping
         self._finalize_distance_metric(layer_to_metrics_mapping)
@@ -244,8 +252,9 @@ class MixedPrecisionSearchManager:
                 else:
                     cfg = self.min_ru_config.copy()
                     cfg[node] = candidate_idx
-                    real_cfg = self.config_reconstruction_helper.reconstruct_full_configuration(cfg)
-                    candidate_rus = self.ru_helper.compute_utilization(self.ru_targets, real_cfg)
+                    if self.using_virtual_graph:
+                        cfg = self.config_reconstructor.reconstruct_full_configuration(cfg)
+                    candidate_rus = self.orig_graph_ru_helper.compute_utilization(self.ru_targets, cfg)
 
                 for target, ru in candidate_rus.items():
                     rus_per_candidate[target].append(ru)
@@ -283,8 +292,8 @@ class MixedPrecisionSearchManager:
         with the given config.
 
         """
-        act_qcs, w_qcs = self.ru_helper.get_quantization_candidates(config)
-        ru = self.ru_helper.ru_calculator.compute_resource_utilization(
+        act_qcs, w_qcs = self.orig_graph_ru_helper.get_quantization_candidates(config)
+        ru = self.orig_graph_ru_helper.ru_calculator.compute_resource_utilization(
             target_criterion=TargetInclusionCriterion.AnyQuantized, bitwidth_mode=BitwidthMode.QCustom, act_qcs=act_qcs,
             w_qcs=w_qcs, ru_targets=self.ru_targets, allow_unused_qcs=True)
         return ru
@@ -303,7 +312,7 @@ class MixedPrecisionSearchManager:
         # normalize metric for numerical stability
         max_dist = max(itertools.chain.from_iterable(layer_to_metrics_mapping.values()))
 
-        if max_dist >= self.sensitivity_evaluator.quant_config.metric_normalization_threshold:
+        if max_dist >= self.mp_config.metric_normalization_threshold:
             Logger.warning(f"The mixed precision distance metric values indicate a large error in the quantized model."
                            f"this can cause numerical issues."
                            f"The program will proceed with mixed precision search after scaling the metric values,"
@@ -387,7 +396,9 @@ class ConfigReconstructionHelper:
 
         return orig_cfg
 
-    def reconstruct_separate_aw_configs(self, virtual_cfg: Dict[BaseNode, int], include_non_configurable: bool) \
+    def reconstruct_separate_aw_configs(self,
+                                        virtual_cfg: Dict[BaseNode, int],
+                                        include_non_configurable: bool = False) \
             -> Tuple[Dict[BaseNode, int], Dict[BaseNode, int]]:
         """
         Retrieves original activation and weights nodes and corresponding candidates for a given configuration of the