mct-nightly 2.2.0.20250115.152408__py3-none-any.whl → 2.2.0.20250117.527__py3-none-any.whl

{mct_nightly-2.2.0.20250115.152408.dist-info → mct_nightly-2.2.0.20250117.527.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: mct-nightly
-Version: 2.2.0.20250115.152408
+Version: 2.2.0.20250117.527
 Summary: A Model Compression Toolkit for neural networks
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: Apache Software License

{mct_nightly-2.2.0.20250115.152408.dist-info → mct_nightly-2.2.0.20250117.527.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=XhsbE4HIyMgv8e1ehZhkkE08uFoS_LDP9ZhVECCHkNM,1557
+model_compression_toolkit/__init__.py,sha256=9wy-eBj_iVmaSe9zp5-Pq8QOeeSZS_srcqrDNVuDsuE,1557
 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
@@ -41,11 +41,11 @@ model_compression_toolkit/core/common/graph/graph_searches.py,sha256=2oKuW6L8hP-
 model_compression_toolkit/core/common/graph/virtual_activation_weights_node.py,sha256=3el-A7j1oyoo1_9zq3faQp7IeRsFXFCvnrb3zZFXpU0,9803
 model_compression_toolkit/core/common/graph/memory_graph/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/core/common/graph/memory_graph/bipartite_graph.py,sha256=X6FK3C3y8ixFRPjC_wm3ClloCX8_06SOdA1TRi7o_LA,3800
-model_compression_toolkit/core/common/graph/memory_graph/compute_graph_max_cut.py,sha256=Na1lAoCJCSQw7XGYsV5xCZg762lbP6Y_uAhsUeLP0yM,2870
-model_compression_toolkit/core/common/graph/memory_graph/cut.py,sha256=GctdLnhsPJgY6UGwRcLNpKE8OLkfVWT3wgby2r9QDD4,2645
-model_compression_toolkit/core/common/graph/memory_graph/max_cut_astar.py,sha256=ZBFIOBBRHuRsiEW31EMwCVb9J7dJo5XBShA_9nnkrRI,17521
+model_compression_toolkit/core/common/graph/memory_graph/compute_graph_max_cut.py,sha256=S3m34BY9P8NPx1I4d9G94X1Zk93MobX5SOVmqipwCOE,3458
+model_compression_toolkit/core/common/graph/memory_graph/cut.py,sha256=7Dfq4TVJIrnencHLJqjhxYKhY7ooUo_ml33WH2IIAgc,2576
+model_compression_toolkit/core/common/graph/memory_graph/max_cut_astar.py,sha256=-Gt4MTnQiyljQWtqMmYKtFKvtZBpj5cPH-Mf6n5Iimo,17753
 model_compression_toolkit/core/common/graph/memory_graph/memory_element.py,sha256=ISD2BvJWj5mB91jrFjG8VQb0oOoLBoita_thCZWzCPI,4238
-model_compression_toolkit/core/common/graph/memory_graph/memory_graph.py,sha256=3OC8kMXuzBv-R7wWmKY-i1AQNAr5x3LBZ4aj7hHF-cQ,7791
+model_compression_toolkit/core/common/graph/memory_graph/memory_graph.py,sha256=FCzK4HmX4lWI4qGoGv94wpGv7o6_f5wPBfeBPMerZ18,7752
 model_compression_toolkit/core/common/hessian/__init__.py,sha256=E7LK3K_1AwMCQokanNc1JODMwUKNOKmwXQiGQ7GO10I,1033
 model_compression_toolkit/core/common/hessian/hessian_info_service.py,sha256=YynbVHdHH2gPlk1QHXH6GygIkXRZ9qxR14cpgKrHPT0,13238
 model_compression_toolkit/core/common/hessian/hessian_info_utils.py,sha256=1axmN0tjJSo_7hUr2d2KMv4y1pBi19cqWSQpi4BbdsA,1458
@@ -73,10 +73,10 @@ model_compression_toolkit/core/common/mixed_precision/set_layer_to_bitwidth.py,s
 model_compression_toolkit/core/common/mixed_precision/solution_refinement_procedure.py,sha256=UWgxzhKWFOoESLq0TFVz0M1PhkU9d9n6wccSA3RgUxk,7903
 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=T5yVr7lay-6QLuTDBZNI1Ufj02EMBWuY_yHjC8eHx5I,3998
-model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_calculator.py,sha256=qkYrYORLL5wmdmCkEY3tDSgabsGYt3OaTDVsgHWYBfE,34885
+model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_calculator.py,sha256=eaZX_Sng1uBpqjKUKuWMQO8wUfnjoQJqEoGwPFD3gsw,35051
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_data.py,sha256=tfcbMppa5KP_brfkFWRiOX9LQVHGXJtlgxyAt9oDGuw,8529
 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=ZvLxFIfMUPAyKKzPhJcuZyjjngLD9_1wWFU8e14vEbA,17176
+model_compression_toolkit/core/common/mixed_precision/search_methods/linear_programming.py,sha256=uhC0az5OVSfeYexcasoy0cT8ZOonFKIedk_1U-ZPLhA,17171
 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
@@ -159,7 +159,7 @@ model_compression_toolkit/core/keras/default_framework_info.py,sha256=IGEHKH3Icm
 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=bbp9jn0pyxcVUkfm_356m-hY2IQUWe_QLz8kclDC7SQ,5453
+model_compression_toolkit/core/keras/resource_utilization_data_facade.py,sha256=XBCmUrHy_fNQCfSjnXCpwuEtc7cda4hXySuiIzhFGqc,5696
 model_compression_toolkit/core/keras/tf_tensor_numpy.py,sha256=jzD8FGEEa8ZD7w8IpTRdp-Udf1MwOTgjg2XTS1Givic,2696
 model_compression_toolkit/core/keras/back2framework/__init__.py,sha256=rhIiXg_nBgUZ-baE3M6SzCuQbcnq4iebY1jtJBvKHOM,808
 model_compression_toolkit/core/keras/back2framework/factory_model_builder.py,sha256=UIQgOOdexycrSKombTMJVvTthR7MlrCihoqM8Kg-rnE,2293
@@ -224,7 +224,7 @@ model_compression_toolkit/core/pytorch/default_framework_info.py,sha256=NLdmiig5
 model_compression_toolkit/core/pytorch/pytorch_device_config.py,sha256=S25cuw10AW3SEN_fRAGRcG_I3wdvvQx1ehSJzPnn-UI,4404
 model_compression_toolkit/core/pytorch/pytorch_implementation.py,sha256=7jzJ4TBKNwwQ9E7W-My8LkmYEJHHNn8weNuO1PCGS10,29830
 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=3whyWyfMIkQYYV-NX6eSyMM2eKpmCnJJ00RqamZouRg,5374
+model_compression_toolkit/core/pytorch/resource_utilization_data_facade.py,sha256=aIHl-dTAC4ISnWSKLD99c-1W3827vfRGyLjMBib-l3s,5618
 model_compression_toolkit/core/pytorch/utils.py,sha256=7VbgcLwtQvdEEc_AJgSOQ3U3KRKCICFPaBirN1fIQxg,3940
 model_compression_toolkit/core/pytorch/back2framework/__init__.py,sha256=H_WixgN0elVWf3exgGYsi58imPoYDj5eYPeh6x4yfug,813
 model_compression_toolkit/core/pytorch/back2framework/factory_model_builder.py,sha256=bwppTPRs6gL96nm7qPiKrNcBj4Krr0yEsOWjRF0aXmQ,2339
@@ -363,7 +363,7 @@ model_compression_toolkit/gptq/keras/gptq_keras_implementation.py,sha256=axBwnCS
 model_compression_toolkit/gptq/keras/gptq_loss.py,sha256=2hzWzsbuVd5XcL85NM57YeOyHxRY0qMArKn8NvQ1UWw,7643
 model_compression_toolkit/gptq/keras/gptq_training.py,sha256=km9tcuugOkRvprGXQZrsq_GPtA3-7Du_-rnbR_Gyups,23228
 model_compression_toolkit/gptq/keras/graph_info.py,sha256=zwoeHX67nJJ5-zYLjzvMXS9TLsy9BsizARbZiDVjVSA,4473
-model_compression_toolkit/gptq/keras/quantization_facade.py,sha256=_3SM3aKJrSayArnOXVu8F5-XCsVmBzjNYHz9-3qRj4E,18534
+model_compression_toolkit/gptq/keras/quantization_facade.py,sha256=onQSR1YPjQ6IZdqzeeqFMs3IeBT-nWLbI0yXuOkdpKs,18827
 model_compression_toolkit/gptq/keras/quantizer/__init__.py,sha256=-DK1CDXvlsnEbki4lukZLpl6Xrbo91_jcqxXlG5Eg6Q,963
 model_compression_toolkit/gptq/keras/quantizer/base_keras_gptq_quantizer.py,sha256=Rbl9urzkmACvVxICSEyJ02qFOBxWK0UQWtysFJzBVZw,4899
 model_compression_toolkit/gptq/keras/quantizer/quant_utils.py,sha256=Vt7Qb8i4JsE4sFtcjpfM4FTXTtfV1t6SwfoNH8a_Iaw,5055
@@ -379,7 +379,7 @@ model_compression_toolkit/gptq/pytorch/gptq_loss.py,sha256=_07Zx_43bnNokwR5S8phI
 model_compression_toolkit/gptq/pytorch/gptq_pytorch_implementation.py,sha256=tECPTavxn8EEwgLaP2zvxdJH6Vg9jC0YOIMJ7857Sdc,1268
 model_compression_toolkit/gptq/pytorch/gptq_training.py,sha256=WtehnyiYXdUXf8-uNpV0mdsalF7YF7eKnL7tcFrzZoE,19549
 model_compression_toolkit/gptq/pytorch/graph_info.py,sha256=4mVM-VvnBaA64ACVdOe6wTGHdMSa2UTLIUe7nACLcdo,4008
-model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=JuMzdeAaS2Ak2NdULsJpOoKju_Kv5L690-ftabr6quo,16631
+model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=Dkanqdv7Eo5lWRoa56aomU5VdH9yqA6zd8I4WE37hxk,16874
 model_compression_toolkit/gptq/pytorch/quantizer/__init__.py,sha256=ZHNHo1yzye44m9_ht4UUZfTpK01RiVR3Tr74-vtnOGI,968
 model_compression_toolkit/gptq/pytorch/quantizer/base_pytorch_gptq_quantizer.py,sha256=fKg-PNOhGBiL-4eySS9Fyw0GkA76Pq8jT_HbJuJ8iZU,4143
 model_compression_toolkit/gptq/pytorch/quantizer/quant_utils.py,sha256=OocYYRqvl7rZ37QT0hTzfJnWGiNCPskg7cziTlR7TRk,3893
@@ -392,20 +392,20 @@ model_compression_toolkit/gptq/pytorch/quantizer/ste_rounding/__init__.py,sha256
 model_compression_toolkit/gptq/pytorch/quantizer/ste_rounding/symmetric_ste.py,sha256=DOlLc4C05TTQN0hZ7xRuqV6wgGp9r2xq7JYun_Hi5jM,8712
 model_compression_toolkit/pruning/__init__.py,sha256=lQMZS8G0pvR1LVi53nnJHNXgLNTan_MWMdwsVxhjrow,1106
 model_compression_toolkit/pruning/keras/__init__.py,sha256=3Lkr37Exk9u8811hw8hVqkGcbTQGcLjd3LLuLC3fa_E,698
-model_compression_toolkit/pruning/keras/pruning_facade.py,sha256=gc00ebAnJEygRETXPxnjfUYE6Ze8zWKVpduhjD0APLs,9072
+model_compression_toolkit/pruning/keras/pruning_facade.py,sha256=-cFNawfLeH0VxYVsauByTvjajt1uiycrkBQ0xcWHQEg,9350
 model_compression_toolkit/pruning/pytorch/__init__.py,sha256=pKAdbTCFM_2BrZXUtTIw0ouKotrWwUDF_hP3rPwCM2k,696
-model_compression_toolkit/pruning/pytorch/pruning_facade.py,sha256=dPGN78I8ZQdcCj_R3DB1hszUJmyRvKFEzZehxjZk-Ro,9757
+model_compression_toolkit/pruning/pytorch/pruning_facade.py,sha256=FmUQvT0T247XaLv8Y6AxBv1G3fCgvndmP1RQdiE3pSU,10044
 model_compression_toolkit/ptq/__init__.py,sha256=Z_hkmTh7aLFei1DJKV0oNVUbrv_Q_0CTw-qD85Xf8UM,904
 model_compression_toolkit/ptq/runner.py,sha256=_c1dSjlPPpsx59Vbg1buhG9bZq__OORz1VlPkwjJzoc,2552
 model_compression_toolkit/ptq/keras/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/ptq/keras/quantization_facade.py,sha256=EeigR6O00Ir4X8nB_T3KsKE939Pg2lbQf5S3VA0orPE,11336
+model_compression_toolkit/ptq/keras/quantization_facade.py,sha256=f8sa46eUNHmeaVs3huhZv14DHm5j1X-VInCYdI7nXAY,11567
 model_compression_toolkit/ptq/pytorch/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/ptq/pytorch/quantization_facade.py,sha256=oqXBlHee7L10heMP11WkiFCNgVAc6RqmDa2HZFWGK0U,9771
+model_compression_toolkit/ptq/pytorch/quantization_facade.py,sha256=p5FwojKaybYdsOUVI7qBNa7R8Nge3EXdu38Jf2jHr84,10021
 model_compression_toolkit/qat/__init__.py,sha256=AaC4KBha4jDW_tyg2SOxZaKh_idIz0gZtDK3_zxs64E,1241
 model_compression_toolkit/qat/common/__init__.py,sha256=6tLZ4R4pYP6QVztLVQC_jik2nES3l4uhML0qUxZrezk,829
 model_compression_toolkit/qat/common/qat_config.py,sha256=xtfVSoyELGXynHNrw86dB9FU3Inu0zwehc3wLrh7JvY,2918
 model_compression_toolkit/qat/keras/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/qat/keras/quantization_facade.py,sha256=cmFRLBVsyv-fyYzz3YY0y1opztcDvDF08Dj9tbHzWvc,17626
+model_compression_toolkit/qat/keras/quantization_facade.py,sha256=DsFAl4FtnR4QQoztUAMD1FgL6DgcdK5jdTp0lk9MHLY,17793
 model_compression_toolkit/qat/keras/quantizer/__init__.py,sha256=zmYyCa25_KLCSUCGUDRslh3RCIjcRMxc_oXa54Aui-4,996
 model_compression_toolkit/qat/keras/quantizer/base_keras_qat_weight_quantizer.py,sha256=EbIt4lMlh6cU4awFLMBp0IlZ2zUUp-WtnlW5Wn19FDM,1793
 model_compression_toolkit/qat/keras/quantizer/quant_utils.py,sha256=cBULOgWUodcBO1lHevZggdTevuDYI6tQceV86U2x6DA,2543
@@ -417,7 +417,7 @@ model_compression_toolkit/qat/keras/quantizer/ste_rounding/__init__.py,sha256=cc
 model_compression_toolkit/qat/keras/quantizer/ste_rounding/symmetric_ste.py,sha256=lXeMPI-n24jbZDGrtOs5eQZ14QvmhFd0e7Y1_QRQxw0,8214
 model_compression_toolkit/qat/keras/quantizer/ste_rounding/uniform_ste.py,sha256=ZdZwMwLa1Ws2eo3DiQYYTvPS1JfiswZL1xlQPtRnIgE,7067
 model_compression_toolkit/qat/pytorch/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/qat/pytorch/quantization_facade.py,sha256=yhMxF3Ah4PRjaTBlmVoRmoCX-pZ0mC9Bq9uslIo6Ud0,13780
+model_compression_toolkit/qat/pytorch/quantization_facade.py,sha256=BvKYsLXyWvE3MXN7khYhBQXVLm-r-C17XpJkEwit7KM,14095
 model_compression_toolkit/qat/pytorch/quantizer/__init__.py,sha256=xYa4C8pr9cG1f3mQQcBXO_u3IdJN-zl7leZxuXDs86w,1003
 model_compression_toolkit/qat/pytorch/quantizer/base_pytorch_qat_weight_quantizer.py,sha256=gjzrnBAZr5c_OrDpSjxpQYa_jKImv7ll52cng07_2oE,1813
 model_compression_toolkit/qat/pytorch/quantizer/quantization_builder.py,sha256=lM10cGUkkTDtRyLLdWj5Rk0cgvcxp0uaCseyvrnk_Vg,5752
@@ -430,7 +430,7 @@ model_compression_toolkit/qat/pytorch/quantizer/ste_rounding/uniform_ste.py,sha2
 model_compression_toolkit/target_platform_capabilities/__init__.py,sha256=9ZcT9JVlYzy8k7MlAXhj086gn6SxlGFsjMvy7ubcnfc,1392
 model_compression_toolkit/target_platform_capabilities/constants.py,sha256=BFSgDwYWU1sZShjoW2S7eH3AI0D4SqDOeOu_sQ971LE,1518
 model_compression_toolkit/target_platform_capabilities/immutable.py,sha256=YhROBiXEIB3TU-bAFrnL3qbAsb1yuWPBAQ_CLOJbYUU,1827
-model_compression_toolkit/target_platform_capabilities/tpc_io_handler.py,sha256=tbDBSSh7sJejDPOfLZ-riGnDfhPqBeIY4ZXqZjZd_eM,4136
+model_compression_toolkit/target_platform_capabilities/tpc_io_handler.py,sha256=4ydTWWKv_PEOAFok2JtxFNj8rav-0IlqcXKF6lnhHNE,4157
 model_compression_toolkit/target_platform_capabilities/schema/__init__.py,sha256=pKAdbTCFM_2BrZXUtTIw0ouKotrWwUDF_hP3rPwCM2k,696
 model_compression_toolkit/target_platform_capabilities/schema/mct_current_schema.py,sha256=PvO8eHxnb3A55gyExT5fZGnOUl3ce7BbbT5SPxCEXNo,541
 model_compression_toolkit/target_platform_capabilities/schema/schema_functions.py,sha256=vBkXxVJagm9JKB9cdm4Pvi7u_luriXUjvNn0-m8Zr0k,4653
@@ -523,8 +523,8 @@ 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.2.0.20250115.152408.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.2.0.20250115.152408.dist-info/METADATA,sha256=bRUgbaQMx5oDrEiJyTqJodZ-mKXrnM38lHCVuiWvSxA,26604
-mct_nightly-2.2.0.20250115.152408.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-mct_nightly-2.2.0.20250115.152408.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.2.0.20250115.152408.dist-info/RECORD,,
+mct_nightly-2.2.0.20250117.527.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.2.0.20250117.527.dist-info/METADATA,sha256=ywy9ErTqUzchvEY5i9iwgRiAi2lr186UmPlZ59OADA4,26601
+mct_nightly-2.2.0.20250117.527.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+mct_nightly-2.2.0.20250117.527.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.2.0.20250117.527.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.20250115.152408"
+__version__ = "2.2.0.20250117.000527"

model_compression_toolkit/core/common/graph/memory_graph/compute_graph_max_cut.py

@@ -13,9 +13,9 @@
 # limitations under the License.
 # ==============================================================================
 from collections import namedtuple
-
 from typing import Tuple, List
 
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import OPERATORS_SCHEDULING, MAX_CUT, CUTS, FUSED_NODES_MAPPING
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.graph.memory_graph.cut import Cut
@@ -49,7 +49,17 @@ def compute_graph_max_cut(memory_graph: MemoryGraph,
     it = 0
     while it < n_iter:
         estimate = (u_bound + l_bound) / 2
-        schedule, max_cut_size, cuts = max_cut_astar.solve(estimate=estimate, iter_limit=astar_n_iter)
+        # Add a timeout of 5 minutes to the solver from the 2nd iteration.
+        try:
+            schedule, max_cut_size, cuts = max_cut_astar.solve(estimate=estimate, iter_limit=astar_n_iter,
+                                                               time_limit=None if it == 0 else 300)
+        except TimeoutError:
+            if last_result[0] is None:
+                Logger.critical(f"Max-cut solver stopped on timeout in iteration {it} before finding a solution.")  # pragma: no cover
+            else:
+                Logger.warning(f"Max-cut solver stopped on timeout in iteration {it}.")
+                return last_result
+
         if schedule is None:
             l_bound = estimate
         else:

model_compression_toolkit/core/common/graph/memory_graph/cut.py

@@ -67,7 +67,7 @@ class Cut:
         return False  # pragma: no cover
 
     def __hash__(self):
-        return hash((frozenset(self.op_order), frozenset(self.op_record), self.mem_elements))
+        return id(self)
 
     def __repr__(self):
-        return f"<Cut: Nodes={[e.node_name for e in self.mem_elements.elements]}, size={self.memory_size()}>"  # pragma: no cover
+        return f"<Cut: Nodes={[e.node_name for e in self.mem_elements.elements]}, size={self.memory_size()}>"  # pragma: no cover

model_compression_toolkit/core/common/graph/memory_graph/max_cut_astar.py

@@ -13,7 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 import copy
-from typing import List, Tuple, Dict
+from typing import List, Tuple, Dict, Set
+from time import time
 
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.constants import DUMMY_TENSOR, DUMMY_NODE
@@ -122,7 +123,7 @@ class MaxCutAstar:
         self.target_cut = Cut([], set(), MemoryElements(elements={target_dummy_b, target_dummy_b2},
                                                         total_size=0))
 
-    def solve(self, estimate: float, iter_limit: int = 500) -> Tuple[List[BaseNode], float, List[Cut]]:
+    def solve(self, estimate: float, iter_limit: int = 500, time_limit: int = None) -> Tuple[List[BaseNode], float, List[Cut]]:
         """
         The AStar solver function. This method runs an AStar-like search on the memory graph,
         using the given estimate as a heuristic gap for solutions to consider.
@@ -131,6 +132,7 @@ class MaxCutAstar:
             estimate: Cut size estimation to consider larger size of nodes in each
                 expansion step, in order to fasten the algorithm divergence towards a solution.
             iter_limit: An upper limit for the number of expansion steps that the algorithm preforms.
+            time_limit: Optional time limit to the solver. Defaults to None which means no limit.
 
         Returns: A solution (if found within the steps limit) which contains:
         - A schedule for computation of the model (List of nodes).
@@ -139,14 +141,17 @@ class MaxCutAstar:
 
         """
 
-        open_list = [self.src_cut]
-        closed_list = []
+        open_list = {self.src_cut}
+        closed_list = set()
         costs = {self.src_cut: self.src_cut.memory_size()}
         routes = {self.src_cut: [self.src_cut]}
 
         expansion_count = 0
 
+        t1 = time()
         while expansion_count < iter_limit and len(open_list) > 0:
+            if time_limit is not None and time() - t1 > time_limit:
+                raise TimeoutError
             # Choose next node to expand
             next_cut = self._get_cut_to_expand(open_list, costs, routes, estimate)
 
@@ -159,22 +164,21 @@ class MaxCutAstar:
 
             if self.is_pivot(next_cut):
                 # Can clear all search history
-                open_list = []
-                closed_list = []
+                open_list.clear()
+                closed_list.clear()
                 routes = {}
             else:
                 # Can remove only next_cut and put it in closed_list
                 open_list.remove(next_cut)
                 del routes[next_cut]
-                closed_list.append(next_cut)
+                closed_list.add(next_cut)
 
             # Expand the chosen cut
             expanded_cuts = self.expand(next_cut)
             expansion_count += 1
 
             # Only consider nodes that where not already visited
-            expanded_cuts = [_c for _c in expanded_cuts if _c not in closed_list]
-            for c in expanded_cuts:
+            for c in filter(lambda _c: _c not in closed_list, expanded_cuts):
                 cost = self.accumulate(cut_cost, c.memory_size())
                 if c not in open_list:
                     self._update_expanded_node(c, cost, cut_route, open_list, costs, routes)
@@ -192,7 +196,7 @@ class MaxCutAstar:
         return None, 0, None  # pragma: no cover
 
     @staticmethod
-    def _update_expanded_node(cut: Cut, cost: float, route: List[Cut], open_list: List[Cut],
+    def _update_expanded_node(cut: Cut, cost: float, route: List[Cut], open_list: Set[Cut],
                               costs: Dict[Cut, float], routes: Dict[Cut, List[Cut]]):
         """
         An auxiliary method for updating search data structures according to an expanded node.
@@ -201,16 +205,16 @@ class MaxCutAstar:
             cut: A cut to expand the search to.
             cost: The cost of the cut.
             route: The rout to the cut.
-            open_list: The search open list.
+            open_list: The search open set.
             costs: The search utility mapping between cuts and their cost.
             routes: The search utility mapping between cuts and their routes.
 
         """
-        open_list.append(cut)
+        open_list.add(cut)
         costs.update({cut: cost})
         routes.update({cut: [cut] + route})
 
-    def _get_cut_to_expand(self, open_list: List[Cut], costs: Dict[Cut, float], routes: Dict[Cut, List[Cut]],
+    def _get_cut_to_expand(self, open_list: Set[Cut], costs: Dict[Cut, float], routes: Dict[Cut, List[Cut]],
                            estimate: float) -> Cut:
         """
         An auxiliary method for finding a cut for expanding the search out of a set of potential cuts for expansion.

model_compression_toolkit/core/common/graph/memory_graph/memory_graph.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 from typing import List
 from operator import getitem
+from functools import cache
 
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.graph.edge import EDGE_SOURCE_INDEX
@@ -82,7 +83,6 @@ class MemoryGraph(DirectedBipartiteGraph):
         inputs_tensors_memory = [sum([t.total_size for t in self.operation_node_children(n)])
                                  for n in nodes if n in model_graph.get_inputs()]
 
-        # TODO maxcut: why both inputs and outputs of each nodes, while the A* solves for node outputs only???
         nodes_total_memory = [sum([t.total_size for t in self.operation_node_children(n)] +
                                   [t.total_size for t in self.operation_node_parents(n)])
                               for n in nodes if n not in model_graph.get_inputs()]
@@ -117,6 +117,7 @@ class MemoryGraph(DirectedBipartiteGraph):
         """
         self.sinks_b = [n for n in self.b_nodes if len(list(self.successors(n))) == 0]
 
+    @cache
     def activation_tensor_children(self, activation_tensor: ActivationMemoryTensor) -> List[BaseNode]:
         """
         Returns the children nodes of a side B node (activation tensor) in the bipartite graph.
@@ -129,6 +130,7 @@ class MemoryGraph(DirectedBipartiteGraph):
         """
         return [oe[1] for oe in self.out_edges(activation_tensor)]
 
+    @cache
     def activation_tensor_parents(self, activation_tensor: ActivationMemoryTensor) -> List[BaseNode]:
         """
         Returns the parents nodes of a side B node (activation tensor) in the bipartite graph.
@@ -141,6 +143,7 @@ class MemoryGraph(DirectedBipartiteGraph):
         """
         return [ie[0] for ie in self.in_edges(activation_tensor)]
 
+    @cache
     def operation_node_children(self, op_node: BaseNode) -> List[ActivationMemoryTensor]:
         """
         Returns the children nodes of a side A node (operation) in the bipartite graph.
@@ -153,6 +156,7 @@ class MemoryGraph(DirectedBipartiteGraph):
         """
         return [oe[1] for oe in self.out_edges(op_node)]
 
+    @cache
     def operation_node_parents(self, op_node: BaseNode) -> List[ActivationMemoryTensor]:
         """
         Returns the parents nodes of a side A node (operation) in the bipartite graph.

model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_calculator.py

@@ -17,6 +17,7 @@ from copy import deepcopy
 from enum import Enum, auto
 from typing import Dict, NamedTuple, Optional, Tuple, List, Iterable, Union, Literal, Sequence, Set
 
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import FLOAT_BITWIDTH
 from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core.common import Graph, BaseNode
@@ -169,6 +170,7 @@ class ResourceUtilizationCalculator:
             w_total, *_ = self.compute_weights_utilization(target_criterion, bitwidth_mode, w_qcs)
 
         if {RUTarget.ACTIVATION, RUTarget.TOTAL}.intersection(ru_targets):
+            Logger.warning("Using an experimental feature max-cut for activation memory utilization estimation.")
             a_total = self.compute_activations_utilization(target_criterion, bitwidth_mode, act_qcs)
 
         ru = ResourceUtilization()

model_compression_toolkit/core/common/mixed_precision/search_methods/linear_programming.py

@@ -16,7 +16,7 @@
 import numpy as np
 from pulp import *
 from tqdm import tqdm
-from typing import Dict, Tuple, Set, Any
+from typing import Dict, Tuple, Any, Optional
 
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization, RUTarget
@@ -182,7 +182,7 @@ def _add_ru_constraints(search_manager: MixedPrecisionSearchManager,
                         target_resource_utilization: ResourceUtilization,
                         indicators_matrix: np.ndarray,
                         lp_problem: LpProblem,
-                        non_conf_ru_dict: Optional[Dict[RUTarget, np.ndarray]]):
+                        non_conf_ru_dict: Dict[RUTarget, np.ndarray]):
     """
     Adding targets constraints for the Lp problem for the given target resource utilization.
     The update to the Lp problem object is done inplace.

model_compression_toolkit/core/keras/resource_utilization_data_facade.py

@@ -13,13 +13,14 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Callable
+from typing import Callable, Union
 from model_compression_toolkit.core import MixedPrecisionQuantizationConfig, CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import TENSORFLOW
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization_data import compute_resource_utilization_data
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TF
 
 if FOUND_TF:
@@ -38,7 +39,7 @@ if FOUND_TF:
                                         representative_data_gen: Callable,
                                         core_config: CoreConfig = CoreConfig(
                                             mixed_precision_config=MixedPrecisionQuantizationConfig()),
-                                        target_platform_capabilities: TargetPlatformCapabilities = KERAS_DEFAULT_TPC
+                                        target_platform_capabilities: Union[TargetPlatformCapabilities, str] = KERAS_DEFAULT_TPC
                                         ) -> ResourceUtilization:
         """
         Computes resource utilization data that can be used to calculate the desired target resource utilization
@@ -50,7 +51,7 @@ if FOUND_TF:
             in_model (Model): Keras model to quantize.
             representative_data_gen (Callable): Dataset used for calibration.
             core_config (CoreConfig): CoreConfig containing parameters for quantization and mixed precision of how the model should be quantized.
-            target_platform_capabilities (FrameworkQuantizationCapabilities): FrameworkQuantizationCapabilities to optimize the Keras model according to.
+            target_platform_capabilities (Union[TargetPlatformCapabilities, str]): FrameworkQuantizationCapabilities to optimize the Keras model according to.
 
         Returns:
 
@@ -81,6 +82,7 @@ if FOUND_TF:
 
         fw_impl = KerasImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach tpc model to framework
         attach2keras = AttachTpcToKeras()
         target_platform_capabilities = attach2keras.attach(

model_compression_toolkit/core/pytorch/resource_utilization_data_facade.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Callable
+from typing import Callable, Union
 
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import PYTORCH
@@ -23,6 +23,7 @@ from model_compression_toolkit.core.common.mixed_precision.resource_utilization_
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig
 from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 
 if FOUND_TORCH:
@@ -40,7 +41,7 @@ if FOUND_TORCH:
     def pytorch_resource_utilization_data(in_model: Module,
                                           representative_data_gen: Callable,
                                           core_config: CoreConfig = CoreConfig(),
-                                          target_platform_capabilities: TargetPlatformCapabilities= PYTORCH_DEFAULT_TPC
+                                          target_platform_capabilities: Union[TargetPlatformCapabilities, str] = PYTORCH_DEFAULT_TPC
                                           ) -> ResourceUtilization:
         """
         Computes resource utilization data that can be used to calculate the desired target resource utilization for mixed-precision quantization.
@@ -50,7 +51,7 @@ if FOUND_TORCH:
             in_model (Model): PyTorch model to quantize.
             representative_data_gen (Callable): Dataset used for calibration.
             core_config (CoreConfig): CoreConfig containing parameters for quantization and mixed precision
-            target_platform_capabilities (FrameworkQuantizationCapabilities): FrameworkQuantizationCapabilities to optimize the PyTorch model according to.
+            target_platform_capabilities (Union[TargetPlatformCapabilities, str]): FrameworkQuantizationCapabilities to optimize the PyTorch model according to.
 
         Returns:
 
@@ -81,6 +82,7 @@ if FOUND_TORCH:
 
         fw_impl = PytorchImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach tpc model to framework
         attach2pytorch = AttachTpcToPytorch()
         target_platform_capabilities = (

model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -25,6 +25,7 @@ from model_compression_toolkit.constants import TENSORFLOW, ACT_HESSIAN_DEFAULT_
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2keras import \
     AttachTpcToKeras
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig, GPTQHessianScoresConfig, \
@@ -156,7 +157,8 @@ if FOUND_TF:
                                                   gptq_representative_data_gen: Callable = None,
                                                   target_resource_utilization: ResourceUtilization = None,
                                                   core_config: CoreConfig = CoreConfig(),
-                                                  target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC) -> Tuple[Model, UserInformation]:
+                                                  target_platform_capabilities: Union[TargetPlatformCapabilities, str]
+                                                  = DEFAULT_KERAS_TPC) -> Tuple[Model, UserInformation]:
         """
         Quantize a trained Keras model using post-training quantization. The model is quantized using a
         symmetric constraint quantization thresholds (power of two).
@@ -180,7 +182,7 @@ if FOUND_TF:
             gptq_representative_data_gen (Callable): Dataset used for GPTQ training. If None defaults to representative_data_gen
             target_resource_utilization (ResourceUtilization): ResourceUtilization object to limit the search of the mixed-precision configuration as desired.
             core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
-            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
+            target_platform_capabilities (Union[TargetPlatformCapabilities, str]): TargetPlatformCapabilities to optimize the Keras model according to.
 
         Returns:
 
@@ -241,6 +243,7 @@ if FOUND_TF:
 
         fw_impl = GPTQKerasImplemantation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach tpc model to framework
         attach2keras = AttachTpcToKeras()
         framework_platform_capabilities = attach2keras.attach(

model_compression_toolkit/gptq/pytorch/quantization_facade.py

@@ -32,6 +32,7 @@ from model_compression_toolkit.gptq.runner import gptq_runner
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.metadata import create_model_metadata
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 
 
@@ -145,7 +146,7 @@ if FOUND_TORCH:
                                                     core_config: CoreConfig = CoreConfig(),
                                                     gptq_config: GradientPTQConfig = None,
                                                     gptq_representative_data_gen: Callable = None,
-                                                    target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYTORCH_TPC):
+                                                    target_platform_capabilities: Union[TargetPlatformCapabilities, str] = DEFAULT_PYTORCH_TPC):
         """
         Quantize a trained Pytorch module using post-training quantization.
         By default, the module is quantized using a symmetric constraint quantization thresholds
@@ -169,7 +170,7 @@ if FOUND_TORCH:
             core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
             gptq_config (GradientPTQConfig): Configuration for using gptq (e.g. optimizer).
             gptq_representative_data_gen (Callable): Dataset used for GPTQ training. If None defaults to representative_data_gen
-            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the PyTorch model according to.
+            target_platform_capabilities (Union[TargetPlatformCapabilities, str]): TargetPlatformCapabilities to optimize the PyTorch model according to.
 
         Returns:
             A quantized module and information the user may need to handle the quantized module.
@@ -214,6 +215,7 @@ if FOUND_TORCH:
 
         fw_impl = GPTQPytorchImplemantation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach tpc model to framework
         attach2pytorch = AttachTpcToPytorch()
         framework_quantization_capabilities = attach2pytorch.attach(target_platform_capabilities,

model_compression_toolkit/pruning/keras/pruning_facade.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Callable, Tuple
+from typing import Callable, Tuple, Union
 
 from model_compression_toolkit import get_target_platform_capabilities
 from model_compression_toolkit.constants import TENSORFLOW
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.pruning.pruner import Pruner
@@ -43,7 +44,8 @@ if FOUND_TF:
                                    target_resource_utilization: ResourceUtilization,
                                    representative_data_gen: Callable,
                                    pruning_config: PruningConfig = PruningConfig(),
-                                   target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC) -> Tuple[Model, PruningInfo]:
+                                   target_platform_capabilities: Union[TargetPlatformCapabilities, str]
+                                   = DEFAULT_KERAS_TPC) -> Tuple[Model, PruningInfo]:
         """
         Perform structured pruning on a Keras model to meet a specified target resource utilization.
         This function prunes the provided model according to the target resource utilization by grouping and pruning
@@ -61,7 +63,7 @@ if FOUND_TF:
             target_resource_utilization (ResourceUtilization): The target Key Performance Indicators to be achieved through pruning.
             representative_data_gen (Callable): A function to generate representative data for pruning analysis.
             pruning_config (PruningConfig): Configuration settings for the pruning process. Defaults to standard config.
-            target_platform_capabilities (FrameworkQuantizationCapabilities): Platform-specific constraints and capabilities. Defaults to DEFAULT_KERAS_TPC.
+            target_platform_capabilities (Union[TargetPlatformCapabilities, str]): Platform-specific constraints and capabilities. Defaults to DEFAULT_KERAS_TPC.
 
         Returns:
             Tuple[Model, PruningInfo]: A tuple containing the pruned Keras model and associated pruning information.
@@ -112,6 +114,7 @@ if FOUND_TF:
         # Instantiate the Keras framework implementation.
         fw_impl = PruningKerasImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach tpc model to framework
         attach2keras = AttachTpcToKeras()
         target_platform_capabilities = attach2keras.attach(target_platform_capabilities)

model_compression_toolkit/pruning/pytorch/pruning_facade.py

@@ -13,10 +13,11 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Callable, Tuple
+from typing import Callable, Tuple, Union
 from model_compression_toolkit import get_target_platform_capabilities
 from model_compression_toolkit.constants import PYTORCH
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.pruning.pruner import Pruner
@@ -47,7 +48,8 @@ if FOUND_TORCH:
                                      target_resource_utilization: ResourceUtilization,
                                      representative_data_gen: Callable,
                                      pruning_config: PruningConfig = PruningConfig(),
-                                     target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYOTRCH_TPC) -> \
+                                     target_platform_capabilities: Union[TargetPlatformCapabilities, str]
+                                     = DEFAULT_PYOTRCH_TPC) -> \
             Tuple[Module, PruningInfo]:
         """
         Perform structured pruning on a Pytorch model to meet a specified target resource utilization.
@@ -66,7 +68,7 @@ if FOUND_TORCH:
             target_resource_utilization (ResourceUtilization): Key Performance Indicators specifying the pruning targets.
             representative_data_gen (Callable): A function to generate representative data for pruning analysis.
             pruning_config (PruningConfig): Configuration settings for the pruning process. Defaults to standard config.
-            target_platform_capabilities (TargetPlatformCapabilities): Platform-specific constraints and capabilities.
+            target_platform_capabilities (Union[TargetPlatformCapabilities, str]): Platform-specific constraints and capabilities.
                 Defaults to DEFAULT_PYTORCH_TPC.
 
         Returns:
@@ -118,6 +120,7 @@ if FOUND_TORCH:
         # Instantiate the Pytorch framework implementation.
         fw_impl = PruningPytorchImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach TPC to framework
         attach2pytorch = AttachTpcToPytorch()
         framework_platform_capabilities = attach2pytorch.attach(target_platform_capabilities)

model_compression_toolkit/ptq/keras/quantization_facade.py

@@ -23,6 +23,7 @@ from model_compression_toolkit.core.common.visualization.tensorboard_writer impo
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import TENSORFLOW
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
@@ -70,7 +71,7 @@ if FOUND_TF:
              representative_data_gen (Callable): Dataset used for calibration.
              target_resource_utilization (ResourceUtilization): ResourceUtilization object to limit the search of the mixed-precision configuration as desired.
              core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
-             target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
+             target_platform_capabilities (Union[TargetPlatformCapabilities, str]): TargetPlatformCapabilities to optimize the Keras model according to.
 
          Returns:
 
@@ -137,6 +138,7 @@ if FOUND_TF:
 
         fw_impl = KerasImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         attach2keras = AttachTpcToKeras()
         framework_platform_capabilities = attach2keras.attach(
             target_platform_capabilities,

model_compression_toolkit/ptq/pytorch/quantization_facade.py

@@ -14,12 +14,13 @@
 # ==============================================================================
 import copy
 
-from typing import Callable
+from typing import Callable, Union
 
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import PYTORCH
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core import CoreConfig
@@ -48,7 +49,7 @@ if FOUND_TORCH:
                                            representative_data_gen: Callable,
                                            target_resource_utilization: ResourceUtilization = None,
                                            core_config: CoreConfig = CoreConfig(),
-                                           target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYTORCH_TPC):
+                                           target_platform_capabilities: Union[TargetPlatformCapabilities, str] = DEFAULT_PYTORCH_TPC):
         """
         Quantize a trained Pytorch module using post-training quantization.
         By default, the module is quantized using a symmetric constraint quantization thresholds
@@ -67,7 +68,7 @@ if FOUND_TORCH:
             representative_data_gen (Callable): Dataset used for calibration.
             target_resource_utilization (ResourceUtilization): ResourceUtilization object to limit the search of the mixed-precision configuration as desired.
             core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
-            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the PyTorch model according to.
+            target_platform_capabilities (Union[TargetPlatformCapabilities, str]): TargetPlatformCapabilities to optimize the PyTorch model according to.
 
         Returns:
             A quantized module and information the user may need to handle the quantized module.
@@ -109,6 +110,7 @@ if FOUND_TORCH:
 
         fw_impl = PytorchImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach tpc model to framework
         attach2pytorch = AttachTpcToPytorch()
         framework_platform_capabilities = attach2pytorch.attach(target_platform_capabilities,

model_compression_toolkit/qat/keras/quantization_facade.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Callable
+from typing import Callable, Union
 from functools import partial
 
 from model_compression_toolkit.core import CoreConfig
@@ -22,6 +22,7 @@ from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2keras import \
     AttachTpcToKeras
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
@@ -49,7 +50,6 @@ if FOUND_TF:
     from model_compression_toolkit.core import common
     from model_compression_toolkit.core.common import BaseNode
     from model_compression_toolkit.constants import TENSORFLOW
-    from model_compression_toolkit.core.common.framework_info import FrameworkInfo
     from model_compression_toolkit.qat.common.qat_config import is_qat_applicable
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
@@ -92,7 +92,7 @@ if FOUND_TF:
                                                             target_resource_utilization: ResourceUtilization = None,
                                                             core_config: CoreConfig = CoreConfig(),
                                                             qat_config: QATConfig = QATConfig(),
-                                                            target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC):
+                                                            target_platform_capabilities: Union[TargetPlatformCapabilities, str] = DEFAULT_KERAS_TPC):
         """
          Prepare a trained Keras model for quantization aware training. First the model quantization is optimized
          with post-training quantization, then the model layers are wrapped with QuantizeWrappers. The model is
@@ -114,7 +114,7 @@ if FOUND_TF:
              target_resource_utilization (ResourceUtilization): ResourceUtilization object to limit the search of the mixed-precision configuration as desired.
              core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
              qat_config (QATConfig): QAT configuration
-             target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
+             target_platform_capabilities (Union[TargetPlatformCapabilities, str]): TargetPlatformCapabilities to optimize the Keras model according to.
 
          Returns:
 
@@ -188,6 +188,7 @@ if FOUND_TF:
 
         fw_impl = KerasImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         attach2keras = AttachTpcToKeras()
         target_platform_capabilities = attach2keras.attach(
             target_platform_capabilities,

model_compression_toolkit/qat/pytorch/quantization_facade.py

@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Callable
+from typing import Callable, Union
 from functools import partial
 
 from model_compression_toolkit.constants import PYTORCH
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2pytorch import \
     AttachTpcToPytorch
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 
 from model_compression_toolkit.core import CoreConfig
@@ -78,7 +79,8 @@ if FOUND_TORCH:
                                                               target_resource_utilization: ResourceUtilization = None,
                                                               core_config: CoreConfig = CoreConfig(),
                                                               qat_config: QATConfig = QATConfig(),
-                                                              target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYTORCH_TPC):
+                                                              target_platform_capabilities: Union[TargetPlatformCapabilities, str]
+                                                              = DEFAULT_PYTORCH_TPC):
         """
          Prepare a trained Pytorch model for quantization aware training. First the model quantization is optimized
          with post-training quantization, then the model layers are wrapped with QuantizeWrappers. The model is
@@ -100,7 +102,7 @@ if FOUND_TORCH:
              target_resource_utilization (ResourceUtilization): ResourceUtilization object to limit the search of the mixed-precision configuration as desired.
              core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
              qat_config (QATConfig): QAT configuration
-             target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Pytorch model according to.
+             target_platform_capabilities (Union[TargetPlatformCapabilities, str]): TargetPlatformCapabilities to optimize the Pytorch model according to.
 
          Returns:
 
@@ -153,10 +155,11 @@ if FOUND_TORCH:
         tb_w = init_tensorboard_writer(DEFAULT_PYTORCH_INFO)
         fw_impl = PytorchImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach tpc model to framework
         attach2pytorch = AttachTpcToPytorch()
         framework_platform_capabilities = attach2pytorch.attach(target_platform_capabilities,
-                                                             core_config.quantization_config.custom_tpc_opset_to_layer)
+                                                                core_config.quantization_config.custom_tpc_opset_to_layer)
 
         # Ignore hessian scores service as we do not use it here
         tg, bit_widths_config, _, _ = core_runner(in_model=in_model,

model_compression_toolkit/target_platform_capabilities/tpc_io_handler.py

@@ -20,13 +20,13 @@ from model_compression_toolkit.target_platform_capabilities.schema.mct_current_s
 import json
 
 
-def load_target_platform_model(tpc_obj_or_path: Union[TargetPlatformCapabilities, str]) -> TargetPlatformCapabilities:
+def load_target_platform_capabilities(tpc_obj_or_path: Union[TargetPlatformCapabilities, str]) -> TargetPlatformCapabilities:
     """
         Parses the tpc input, which can be either a TargetPlatformCapabilities object
         or a string path to a JSON file.
 
         Parameters:
-            tpc_obj_or_path (Union[TargetPlatformModel, str]): Input target platform model or path to .JSON file.
+            tpc_obj_or_path (Union[TargetPlatformCapabilities, str]): Input target platform model or path to .JSON file.
 
         Returns:
             TargetPlatformCapabilities: The parsed TargetPlatformCapabilities.
@@ -66,7 +66,7 @@ def load_target_platform_model(tpc_obj_or_path: Union[TargetPlatformCapabilities
     )
 
 
-def export_target_platform_model(model: TargetPlatformCapabilities, export_path: Union[str, Path]) -> None:
+def export_target_platform_capabilities(model: TargetPlatformCapabilities, export_path: Union[str, Path]) -> None:
     """
     Exports a TargetPlatformCapabilities instance to a JSON file.