mct-nightly 1.11.0.20240303.post423__py3-none-any.whl → 1.11.0.20240305.post352__py3-none-any.whl

{mct_nightly-1.11.0.20240303.post423.dist-info → mct_nightly-1.11.0.20240305.post352.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 1.11.0.20240303.post423
+Version: 1.11.0.20240305.post352
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN
@@ -173,11 +173,11 @@ This pruning technique is designed to compress models for specific hardware arch
 taking into account the target platform's Single Instruction, Multiple Data (SIMD) capabilities. 
 By pruning groups of channels (SIMD groups), our approach not only reduces model size 
 and complexity, but ensures that better utilization of channels is in line with the SIMD architecture 
-for a target KPI of weights memory footprint.
+for a target KPI of weights memory footprint. 
+[Keras API](https://sony.github.io/model_optimization/docs/api/experimental_api_docs/methods/keras_pruning_experimental.html)
+[Pytorch API](https://sony.github.io/model_optimization/docs/api/experimental_api_docs/methods/pytorch_pruning_experimental.html)
 
 
-<u>_Note: Currently, only Keras models pruning is supported._</u>
-
 #### Results
 
 Results for applying pruning to reduce the parameters of the following models by 50%:
@@ -185,7 +185,7 @@ Results for applying pruning to reduce the parameters of the following models by
 | Model           | Dense Model Accuracy | Pruned Model Accuracy |
 |-----------------|----------------------|-----------------------|
 | ResNet50 [2]    | 75.1                 | 72.4                  |
-| DenseNet121 [2] | 75.0                 | 71.15                 |
+| DenseNet121 [3] | 74.44                | 71.71                 |
 
 
 

{mct_nightly-1.11.0.20240303.post423.dist-info → mct_nightly-1.11.0.20240305.post352.dist-info}/RECORD

@@ -2,7 +2,7 @@ model_compression_toolkit/__init__.py,sha256=7NFeY28X4NSHeLTq1JcKkAnp031EMjorU-I
 model_compression_toolkit/constants.py,sha256=_OW_bUeQmf08Bb4oVZ0KfUt-rcCeNOmdBv3aP7NF5fM,3631
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=b9DVktZ-LymFcRxv2aL_sdiE6S2sSrFGWltx6dgEuUY,4863
-model_compression_toolkit/core/__init__.py,sha256=AXC_DLak39yxBGURNPm6LjtxB1P7iD1O3Iq-q6AA5xM,1874
+model_compression_toolkit/core/__init__.py,sha256=P-7OYR4TFYxVV_ZpIJBogkX8bGvXcijlF65Ez3ivjhc,1838
 model_compression_toolkit/core/analyzer.py,sha256=dbsD61pakp_9JXNyAScLdtJvcXny9jr_cMbET0Bd3Sg,2975
 model_compression_toolkit/core/exporter.py,sha256=U_-ea-zYHsnIt2ydameMLZ_gzDaCMI1dRa5IjA8RUuc,4233
 model_compression_toolkit/core/graph_prep_runner.py,sha256=3xp0WYqyeRdlBkf5R6uD2zWubg_JPttOwS7JRhKykBY,10043
@@ -11,7 +11,7 @@ model_compression_toolkit/core/runner.py,sha256=RgN9l0v7aFYu6MTuIZGAB2syr6NBqG_v
 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/data_loader.py,sha256=jCoVIb4yeOWyCrCNRB1W-mgLSyqNVGEepFXrIqufVc4,4119
-model_compression_toolkit/core/common/framework_implementation.py,sha256=pWbFiioYqL69-_bGNdxerLzGgP-HSiegRzKyGkAo27w,20533
+model_compression_toolkit/core/common/framework_implementation.py,sha256=n3T0uOfeni_P-5ut4ModjPdtRUWCWzfAIOW0LdOlVb4,20529
 model_compression_toolkit/core/common/framework_info.py,sha256=1ZMMGS9ip-kSflqkartyNRt9aQ5ub1WepuTRcTy-YSQ,6337
 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
@@ -62,10 +62,10 @@ model_compression_toolkit/core/common/mixed_precision/bit_width_setter.py,sha256
 model_compression_toolkit/core/common/mixed_precision/configurable_quant_id.py,sha256=LLDguK7afsbN742ucLpmJr5TUfTyFpK1vbf2bpVr1v0,882
 model_compression_toolkit/core/common/mixed_precision/configurable_quantizer_utils.py,sha256=kmyBcqGh3qYqo42gIZzouQEljTNpF9apQt6cXEVkTQ0,3871
 model_compression_toolkit/core/common/mixed_precision/distance_weighting.py,sha256=x0cweemRG3_7FlvAbxFK5Zi77qpoKAGqtGndY8MtgwM,2222
-model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py,sha256=n7ZwJSd-6OT59Zg8eJW9PW_NjHy-2q2F4IKyV-_p9TM,8032
-model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py,sha256=eRohepDyOr_cXVjgp2EnntPSIbWDADH0DIjP2Wc15oQ,7107
+model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py,sha256=wKiG-d_qg5dgPpZRK6dvP0MCXUKIfDNpeCHQrI3JxqY,4245
+model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py,sha256=mN-QeabIu_Mz1IzPeQjqgqprCTdwGm4ThYX0gZAek-E,7103
 model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py,sha256=XrHMQAzIUwyHqrY30cwK1lP5jKI0toeWFEG4e3qKzlY,35988
-model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py,sha256=xo_RyCUf2I-nvjof6p_40QMgiUvNpr77iyTZhH9u8ug,28227
+model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py,sha256=z_0hz6mX27Q-Xzg-ftvoyPrw_i9faBGslkP0u0WGxTg,28223
 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=A0-XsePLlqeYDgAFrDLJjYa4mY3VFKRWYQ3vHpPE-c4,6326
 model_compression_toolkit/core/common/mixed_precision/kpi_tools/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
@@ -83,7 +83,7 @@ model_compression_toolkit/core/common/network_editors/node_filters.py,sha256=uML
 model_compression_toolkit/core/common/pruning/__init__.py,sha256=DGJybkDQtKMSMFoZ-nZ3ZifA8uJ6G_D20wHhKHNlmU0,699
 model_compression_toolkit/core/common/pruning/channels_grouping.py,sha256=4jsr1xEBNpV2c4ipi366IfHoHCJVqoRUTTOJdlRomvc,3892
 model_compression_toolkit/core/common/pruning/greedy_mask_calculator.py,sha256=Go3FPHEyvm_Ra0c3XmdUVEFLCGWFoddMuYXd3Trdxkw,7594
-model_compression_toolkit/core/common/pruning/memory_calculator.py,sha256=3VcvCfJGsYSS0VVdP5uaWoI2cncBYtSEU-PwF-OG3E0,18320
+model_compression_toolkit/core/common/pruning/memory_calculator.py,sha256=RnMmgNDHekKFOj-b-ad5rhjuKUvbVawy1A31nxuCRTg,19217
 model_compression_toolkit/core/common/pruning/prune_graph.py,sha256=ddbZLuWvlNoj5so_5NRbIuG5qDFxD9ApG2gPirbov8o,3317
 model_compression_toolkit/core/common/pruning/pruner.py,sha256=vXxzBXQ-oAEnw6PAD1SUiNXX7Xix4JJ0LAmV04sjFz0,7313
 model_compression_toolkit/core/common/pruning/pruning_config.py,sha256=IfF824hNttyw2i4Tuf3g8CUfelJR3eZuOLzf2aEZNAM,3442
@@ -99,19 +99,18 @@ model_compression_toolkit/core/common/pruning/mask/per_channel_mask.py,sha256=hk
 model_compression_toolkit/core/common/pruning/mask/per_simd_group_mask.py,sha256=gmzD32xsfJH8vkkqaspS7vYa6VWayk1GJe-NfoAEugQ,5901
 model_compression_toolkit/core/common/quantization/__init__.py,sha256=sw7LOPN1bM82o3SkMaklyH0jw-TLGK0-fl2Wq73rffI,697
 model_compression_toolkit/core/common/quantization/candidate_node_quantization_config.py,sha256=VW0vGCPKMxbhl2cB_zHx3g9c7cNS4ctVEAvnaNq17jw,5153
-model_compression_toolkit/core/common/quantization/core_config.py,sha256=8DRM4Ar4Er-bllo56LG-Lcx9U2Ebd3jJctf4t2hOcXc,2021
+model_compression_toolkit/core/common/quantization/core_config.py,sha256=IkD4Jl9PWdPucfUMq0TtyUl5DBJvha7Dd2xSW7_7dz8,2015
 model_compression_toolkit/core/common/quantization/debug_config.py,sha256=HtkMmneN-EmAzgZK4Vp4M8Sqm5QKdrvNyyZMpaVqYzY,1482
 model_compression_toolkit/core/common/quantization/filter_nodes_candidates.py,sha256=4GCr4Z6pRMbxIAnq4s7YtdMSqwbRwUzTzCFfs2ahVfk,6137
 model_compression_toolkit/core/common/quantization/node_quantization_config.py,sha256=Q634XzMtjqReiLni8974y13apzbZ9nref-XBGjH17-0,16761
 model_compression_toolkit/core/common/quantization/quantization_config.py,sha256=0547euhiaEX4vgpvIwHd7-pZ3iI7QmIc6Y_qHV4Y5sY,6713
-model_compression_toolkit/core/common/quantization/quantization_fn_selection.py,sha256=9GWuVMW9ifWnflUjLZoHZtcd8NpHyHTUzCdJFKIaGlo,2352
-model_compression_toolkit/core/common/quantization/quantization_params_fn_selection.py,sha256=sEPDeClFxh0uHEGznX7E3bSOJ_t0kUvyWcdxcyZJdwA,4090
+model_compression_toolkit/core/common/quantization/quantization_fn_selection.py,sha256=f2Qa2majjO-gIN3lxqsA8icKJ9FMP-sKbw3lI6XNgBg,2137
+model_compression_toolkit/core/common/quantization/quantization_params_fn_selection.py,sha256=mrgVzZszWjxnjT8zm77UVLWKTOwd2thGBo6WNqAS4X8,3867
 model_compression_toolkit/core/common/quantization/quantize_graph_weights.py,sha256=_OQEFAdYDTHu2Qp-qs02Z1CDxugUKG6k5eCePS1WpXY,2939
 model_compression_toolkit/core/common/quantization/quantize_node.py,sha256=UK_YshvZI0-LrKeT9gFGYcMA7pma1kaR5JAfzJH3HNw,3614
 model_compression_toolkit/core/common/quantization/set_node_quantization_config.py,sha256=qkrWJXLyDSIJhvT8tO9Nh51f4abyVR8zMFuaaMRRrRw,12304
-model_compression_toolkit/core/common/quantization/quantization_params_generation/__init__.py,sha256=_U4IFPuzGyyAymjDjsPl2NF6UbFggqBaiA1Td3sug3I,1608
+model_compression_toolkit/core/common/quantization/quantization_params_generation/__init__.py,sha256=eCDGwsWYLU6z7qbEVb4TozMW_nd5VEP_iCJ6PcvyEPw,1486
 model_compression_toolkit/core/common/quantization/quantization_params_generation/error_functions.py,sha256=rwCedE0zggamSBY50rqh-xqZpIMrn8o96YH_jMCuPrk,16505
-model_compression_toolkit/core/common/quantization/quantization_params_generation/kmeans_params.py,sha256=qDfJbvY64KLOG6n18ddEPTFGrKHlaXzZ136TrVpgH9s,2917
 model_compression_toolkit/core/common/quantization/quantization_params_generation/lut_kmeans_params.py,sha256=h8Zmpq3KdcsdUUy7K1fvWOVSki0mxT8wtKZXGmgFl74,7405
 model_compression_toolkit/core/common/quantization/quantization_params_generation/outlier_filter.py,sha256=9gnfJV89jpGwAx8ImJ5E9NjCv3lDtbyulP4OtgWb62M,1772
 model_compression_toolkit/core/common/quantization/quantization_params_generation/power_of_two_selection.py,sha256=W4j9IB1Grj_Ku1pLjPxb-HLcYU9LTDuf9_0JilbqU2w,8484
@@ -122,7 +121,6 @@ model_compression_toolkit/core/common/quantization/quantization_params_generatio
 model_compression_toolkit/core/common/quantization/quantization_params_generation/symmetric_selection.py,sha256=53OFL3IZxtH-UPakf3h_LZkaZAa3cgc2oqgMUe3Sg8o,9689
 model_compression_toolkit/core/common/quantization/quantization_params_generation/uniform_selection.py,sha256=oiJn1twYpTaq_z5qX4d8_nnk_jouYWHq8z0WAknl5oE,7879
 model_compression_toolkit/core/common/quantization/quantizers/__init__.py,sha256=mjbqLD-KcG3eNeCYpu1GBS7VclGVOQ63x2p6mAAuba4,698
-model_compression_toolkit/core/common/quantization/quantizers/kmeans_quantizer.py,sha256=ZS3IXGbUTW580vwVF5jgxfPVYVL3tQrpvoSqDxVu7zQ,2325
 model_compression_toolkit/core/common/quantization/quantizers/lut_kmeans_quantizer.py,sha256=P0x_y18LypBxP2tV9OWizheYfILqvaMC8RwHo04sUpQ,2761
 model_compression_toolkit/core/common/quantization/quantizers/quantizers_helpers.py,sha256=5JuPwb9HDHaYQj1YyNWGY7GdjJ105Yr8iEEZhzfuRW4,14190
 model_compression_toolkit/core/common/quantization/quantizers/uniform_quantizers.py,sha256=FVeuK-LeuAsRFcqo5uaNHmb6oTOFs21ltghtqswl6KM,5486
@@ -152,10 +150,10 @@ model_compression_toolkit/core/keras/__init__.py,sha256=mjbqLD-KcG3eNeCYpu1GBS7V
 model_compression_toolkit/core/keras/constants.py,sha256=YhuzRqXAdkRFzLT5lRD_jtLVYcUb-d4fUm-D49z5XOg,3158
 model_compression_toolkit/core/keras/custom_layer_validation.py,sha256=f-b14wuiIgitBe7d0MmofYhDCTO3IhwJgwrh-Hq_t_U,1192
 model_compression_toolkit/core/keras/default_framework_info.py,sha256=Ha4HTHuiw_KTS5Po1Xnv6GyK9eprpDhYWf-eooS62Ys,4961
-model_compression_toolkit/core/keras/keras_implementation.py,sha256=j6sJnH5UUpnyXAji01XVl8UiuHQ7JFCUpm7pRLXwYas,28148
+model_compression_toolkit/core/keras/keras_implementation.py,sha256=kvrZeq3gmHeJKZXPyz55NGwO3EseTgo6022ZCWaUPW8,28144
 model_compression_toolkit/core/keras/keras_model_validation.py,sha256=1wNV2clFdC9BzIELRLSO2uKf0xqjLqlkTJudwtCeaJk,1722
 model_compression_toolkit/core/keras/keras_node_prior_info.py,sha256=Aqh31wOPaiZcJIOm-uJwzev0eTMdJyXaOk97rs4z7BU,3879
-model_compression_toolkit/core/keras/kpi_data_facade.py,sha256=0OkwSeuDMUQRelgvp4mz8enLsaSamsp5ePGkRKrZQOU,4826
+model_compression_toolkit/core/keras/kpi_data_facade.py,sha256=nZdRhQuIXjoL3sq2HjffKHWfnafVM-j_oqk0Cc5Op3I,4889
 model_compression_toolkit/core/keras/tf_tensor_numpy.py,sha256=1kBs9URqZTfmRXAsCqvnekV5bKUL3MyqGbORewLIwu8,2457
 model_compression_toolkit/core/keras/back2framework/__init__.py,sha256=rhIiXg_nBgUZ-baE3M6SzCuQbcnq4iebY1jtJBvKHOM,808
 model_compression_toolkit/core/keras/back2framework/factory_model_builder.py,sha256=GSh1Piz5qpA7IlvHTMqUvPn7WBDa0IHEDZdd_TzY9XA,2226
@@ -192,7 +190,7 @@ model_compression_toolkit/core/keras/mixed_precision/__init__.py,sha256=sw7LOPN1
 model_compression_toolkit/core/keras/mixed_precision/configurable_activation_quantizer.py,sha256=s4m5jR2OsBCgDY0W0rB2T3YQwZPQvhZvNfSAPc9Vggo,4922
 model_compression_toolkit/core/keras/mixed_precision/configurable_weights_quantizer.py,sha256=X_1qI-1iD6HxG6ERe48OlCelw-nn3k4NZpeV1vtfNOs,6363
 model_compression_toolkit/core/keras/pruning/__init__.py,sha256=3Lkr37Exk9u8811hw8hVqkGcbTQGcLjd3LLuLC3fa_E,698
-model_compression_toolkit/core/keras/pruning/pruning_keras_implementation.py,sha256=olRbKncMy159Mp2C9UC8VomQ55rDXYVx2hCJ-atQsic,12455
+model_compression_toolkit/core/keras/pruning/pruning_keras_implementation.py,sha256=KUnSnpx8ty8NmNc6HA-HiWg-FBg7rBwkbLgOM_5ZfJA,12676
 model_compression_toolkit/core/keras/quantizer/__init__.py,sha256=mjbqLD-KcG3eNeCYpu1GBS7VclGVOQ63x2p6mAAuba4,698
 model_compression_toolkit/core/keras/quantizer/base_quantizer.py,sha256=eMRjAUU189-AVwNGMlV0M-ZlL48ZYmILzutheUT00xU,1628
 model_compression_toolkit/core/keras/quantizer/fake_quant_builder.py,sha256=Oi64CD83OopPoQNAarl2MJRbCKujU2W8Wdrs9KOPNWk,6151
@@ -211,11 +209,11 @@ model_compression_toolkit/core/keras/statistics_correction/__init__.py,sha256=9H
 model_compression_toolkit/core/keras/statistics_correction/apply_second_moment_correction.py,sha256=XNCtT9klMcsO1v5KA3MmCq_WgXOIT5QSzbfTOa9T-04,3060
 model_compression_toolkit/core/keras/visualization/__init__.py,sha256=mjbqLD-KcG3eNeCYpu1GBS7VclGVOQ63x2p6mAAuba4,698
 model_compression_toolkit/core/pytorch/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
-model_compression_toolkit/core/pytorch/constants.py,sha256=58SdcH1N3mfU4sAx9E54EdizY04bY4gKSf10dCGbW0U,2534
+model_compression_toolkit/core/pytorch/constants.py,sha256=NI-J7REuxn06oEIHsmJ4GqtNC3TbV8xlkJjt5Ar-c4U,2626
 model_compression_toolkit/core/pytorch/default_framework_info.py,sha256=r1XyzUFvrjGcJHQM5ETLsMZIG2yHCr9HMjqf0ti9inw,4175
-model_compression_toolkit/core/pytorch/kpi_data_facade.py,sha256=etYbVo5dwkBwjt8dqlrpGFijvd9y8oP_KX4JsdubnQM,4978
+model_compression_toolkit/core/pytorch/kpi_data_facade.py,sha256=xQdg8vtkwhx1uzElrh0KtwKdWFr6b2Guzv913iA_VoI,4978
 model_compression_toolkit/core/pytorch/pytorch_device_config.py,sha256=IoMvTch5awAEPvB6Tg6ANhFGXvfSgv7JLsUBlxpMwk4,4330
-model_compression_toolkit/core/pytorch/pytorch_implementation.py,sha256=KGwEgOb4jwSWGi9mGd6DDBLcWOcveq41SMjx_0z-wlE,26078
+model_compression_toolkit/core/pytorch/pytorch_implementation.py,sha256=EEHOnsXsvWlNm4x-4ADuFHsIc8vLrZP9pMh3raegaNg,26074
 model_compression_toolkit/core/pytorch/pytorch_node_prior_info.py,sha256=n_B4a6FMwM9D2w8kzy3oenBWZgXNZuIZgTJC6JEuTy0,3250
 model_compression_toolkit/core/pytorch/utils.py,sha256=Bxep5o-Zw3O_iLdlzidiU186r0O6MXXtkJIZm-L1PiA,2847
 model_compression_toolkit/core/pytorch/back2framework/__init__.py,sha256=H_WixgN0elVWf3exgGYsi58imPoYDj5eYPeh6x4yfug,813
@@ -254,6 +252,8 @@ model_compression_toolkit/core/pytorch/hessian/weights_trace_hessian_calculator_
 model_compression_toolkit/core/pytorch/mixed_precision/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/core/pytorch/mixed_precision/configurable_activation_quantizer.py,sha256=iEtLjZw4l8Oyj81QnxbQFRKVuTbQEI98P5LuU7SNGq8,4656
 model_compression_toolkit/core/pytorch/mixed_precision/configurable_weights_quantizer.py,sha256=VUN9vvWQWAh281C0xgV3w4T2DkSaxFZ-xmBgF50vGdo,5961
+model_compression_toolkit/core/pytorch/pruning/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
+model_compression_toolkit/core/pytorch/pruning/pruning_pytorch_implementation.py,sha256=a1-P0ZI8jO1qlYNWqas9AvRb3bwcrNluul6qPc140R8,14601
 model_compression_toolkit/core/pytorch/quantizer/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/core/pytorch/quantizer/fake_quant_builder.py,sha256=rox-f5wbRyxU1UHeHyaoIDXB9r9fCXm1dPN4FVwHqTc,6464
 model_compression_toolkit/core/pytorch/quantizer/lut_fake_quant.py,sha256=uyeBtNokyDUikk-YkDP_mN_2DX0J5oPm3kSfdSUT2Ck,4420
@@ -341,7 +341,7 @@ model_compression_toolkit/gptq/keras/gptq_keras_implementation.py,sha256=axBwnCS
 model_compression_toolkit/gptq/keras/gptq_loss.py,sha256=rbRkF15MYd6nq4G49kcjb_dPTa-XNq9cTkrb93mXawo,6241
 model_compression_toolkit/gptq/keras/gptq_training.py,sha256=75j56X2AcNI_5hInsLvXnWZOGMZIKQY2hStVKBaA_Bc,17705
 model_compression_toolkit/gptq/keras/graph_info.py,sha256=FIGqzJbG6GkdHenvdMu-tGTjp4j9BewdF_spmWCb4Mo,4627
-model_compression_toolkit/gptq/keras/quantization_facade.py,sha256=UryGNh9JudVsPeUYgXRsC3zqPuWjlQlPxpt3v9YW00w,15170
+model_compression_toolkit/gptq/keras/quantization_facade.py,sha256=STtrf7bzqVr1P1Pah3cftGRcYSF7i-tSkcjrcuMdpTw,15162
 model_compression_toolkit/gptq/keras/quantizer/__init__.py,sha256=-DK1CDXvlsnEbki4lukZLpl6Xrbo91_jcqxXlG5Eg6Q,963
 model_compression_toolkit/gptq/keras/quantizer/base_keras_gptq_quantizer.py,sha256=8NrJBftKFbMAF_jYaAbLP6GBwpCv3Ln1NKURaV75zko,4770
 model_compression_toolkit/gptq/keras/quantizer/quant_utils.py,sha256=Vt7Qb8i4JsE4sFtcjpfM4FTXTtfV1t6SwfoNH8a_Iaw,5055
@@ -358,7 +358,7 @@ model_compression_toolkit/gptq/pytorch/gptq_loss.py,sha256=kDuWw-6zh17wZpYWh4Xa9
 model_compression_toolkit/gptq/pytorch/gptq_pytorch_implementation.py,sha256=tECPTavxn8EEwgLaP2zvxdJH6Vg9jC0YOIMJ7857Sdc,1268
 model_compression_toolkit/gptq/pytorch/gptq_training.py,sha256=RfkJNQbeqGZQvlmV0dZO7YJ894Gx2asLnnIHFdWNEZ0,15078
 model_compression_toolkit/gptq/pytorch/graph_info.py,sha256=-0GDC2cr-XXS7cTFTnDflJivGN7VaPnzVPsxCE-vZNU,3955
-model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=ilP9B4rkSWOB94HF5fR8NKl3mZdYrdDARSxjKW92JvQ,13324
+model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=Mr0AWmqaxdr9O4neBej4hB6BWxwP_okHjQ2SwUoVvrM,13318
 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=Zb-P0yRyZHHBlDvUBdRwxDpdduEJyJp6OT9pfKFF5ks,4171
 model_compression_toolkit/gptq/pytorch/quantizer/quant_utils.py,sha256=OocYYRqvl7rZ37QT0hTzfJnWGiNCPskg7cziTlR7TRk,3893
@@ -370,20 +370,22 @@ model_compression_toolkit/gptq/pytorch/quantizer/soft_rounding/symmetric_soft_qu
 model_compression_toolkit/gptq/pytorch/quantizer/soft_rounding/uniform_soft_quantizer.py,sha256=FgPSKoV8p8y-gLNz359XdOPD6w_wpDvcJFtTNLWqYb0,9099
 model_compression_toolkit/gptq/pytorch/quantizer/ste_rounding/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/gptq/pytorch/quantizer/ste_rounding/symmetric_ste.py,sha256=6uxq_w62jn8DDOt9T7VtA6jZ8jTAPcbTufKFOYpVUm4,8768
-model_compression_toolkit/pruning/__init__.py,sha256=inDsxrBI8aXkYI9ZeIeNVTwa5D4cv2ZQ2q9v7hb0_Og,1008
+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=PHKZYBHrVyR348-a6gw44NrV8Ra9iaeFJ0WbWYpzX8k,8020
+model_compression_toolkit/pruning/pytorch/__init__.py,sha256=pKAdbTCFM_2BrZXUtTIw0ouKotrWwUDF_hP3rPwCM2k,696
+model_compression_toolkit/pruning/pytorch/pruning_facade.py,sha256=hO13kbCIyoAgVW6hRoJCUpHr6ThOuqz1Vkg-cftJY5k,8906
 model_compression_toolkit/ptq/__init__.py,sha256=50QBTXOKdj9XLjXtrvf0mhC9FlW6TOi9-pjl96RLR14,930
 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=39-bHTOZ7OENxvbC_HKelwjSO3e8BccGexgwWPClIDk,9969
+model_compression_toolkit/ptq/keras/quantization_facade.py,sha256=J6luWnjkyk8KvXwZnVb5sEG9fynhXIOBMbIRPXDs2g8,9805
 model_compression_toolkit/ptq/pytorch/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/ptq/pytorch/quantization_facade.py,sha256=2r_2BRAocLiMciGBa59wbpDt6U4RWMLsJW0hIwcIbLQ,8722
+model_compression_toolkit/ptq/pytorch/quantization_facade.py,sha256=gsr_tPfY9HbAxKWaHgiq3J_osXUkOoazdLuvWdyyg3A,8567
 model_compression_toolkit/qat/__init__.py,sha256=BYKgH1NwB9fqF1TszULQ5tDfLI-GqgZV5sao-lDN9EM,1091
 model_compression_toolkit/qat/common/__init__.py,sha256=6tLZ4R4pYP6QVztLVQC_jik2nES3l4uhML0qUxZrezk,829
 model_compression_toolkit/qat/common/qat_config.py,sha256=kbSxFL6_u28furq5mW_75STWDmyX4clPt-seJAnX3IQ,3445
 model_compression_toolkit/qat/keras/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/qat/keras/quantization_facade.py,sha256=hB_VrhSqwOjGnOT8BYpXkh52EMzo7I-62-IEGohg_74,16253
+model_compression_toolkit/qat/keras/quantization_facade.py,sha256=3hdUBm-Q8mRe_sBrpHKFguVaUQPEH95ZeiyYN_EclwI,16073
 model_compression_toolkit/qat/keras/quantizer/__init__.py,sha256=zmYyCa25_KLCSUCGUDRslh3RCIjcRMxc_oXa54Aui-4,996
 model_compression_toolkit/qat/keras/quantizer/base_keras_qat_quantizer.py,sha256=gPuIgQb8OafvC3SuA8jNsGoy8S8eTsDCEKuh36WDNss,2104
 model_compression_toolkit/qat/keras/quantizer/quant_utils.py,sha256=cBULOgWUodcBO1lHevZggdTevuDYI6tQceV86U2x6DA,2543
@@ -395,7 +397,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=I4KlaGv17k71IyjuSG9M0OlXlD5P0pfvKa6oCyRQ5FE,13517
 model_compression_toolkit/qat/keras/quantizer/ste_rounding/uniform_ste.py,sha256=EED6LfqhX_OhDRJ9e4GwbpgNC9vq7hoXyJS2VPvG2qc,10789
 model_compression_toolkit/qat/pytorch/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/qat/pytorch/quantization_facade.py,sha256=Kv2vUKY2cSt6kYB7J2B1SC_PGzSFqecpxDcKWGeRzuQ,12629
+model_compression_toolkit/qat/pytorch/quantization_facade.py,sha256=pNBkkT9Ey1CmrQUqg-NHdf4tFVjgoVldX8bqXJp9lik,12473
 model_compression_toolkit/qat/pytorch/quantizer/__init__.py,sha256=xYa4C8pr9cG1f3mQQcBXO_u3IdJN-zl7leZxuXDs86w,1003
 model_compression_toolkit/qat/pytorch/quantizer/base_pytorch_qat_quantizer.py,sha256=FnhuFCuQoSf78FM1z1UZgXXd3k-mKSM7i9dYOuJUmeA,2213
 model_compression_toolkit/qat/pytorch/quantizer/quantization_builder.py,sha256=GOYRDXvQSGe_iUFVmvDy5BqC952hu_-rQO06n8QCyw0,5491
@@ -470,8 +472,8 @@ model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py,sha
 model_compression_toolkit/trainable_infrastructure/keras/quantizer_utils.py,sha256=MVwXNymmFRB2NXIBx4e2mdJ1RfoHxRPYRgjb1MQP5kY,1797
 model_compression_toolkit/trainable_infrastructure/pytorch/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py,sha256=SbvRlIdE32PEBsINt1bhSqvrKL_zbM9V-aeSkOn-sw4,3083
-mct_nightly-1.11.0.20240303.post423.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-1.11.0.20240303.post423.dist-info/METADATA,sha256=7TlXG0QPNK4ZWwDIaw_T_tOAkcIM7jwcxjA1v6RVudQ,17187
-mct_nightly-1.11.0.20240303.post423.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-mct_nightly-1.11.0.20240303.post423.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-1.11.0.20240303.post423.dist-info/RECORD,,
+mct_nightly-1.11.0.20240305.post352.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-1.11.0.20240305.post352.dist-info/METADATA,sha256=_XbUd27NO_Oa_pHVS9Jw32qfVPNAtdEeCOqXXLAoPhA,17377
+mct_nightly-1.11.0.20240305.post352.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+mct_nightly-1.11.0.20240305.post352.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-1.11.0.20240305.post352.dist-info/RECORD,,

model_compression_toolkit/core/__init__.py

@@ -22,6 +22,6 @@ from model_compression_toolkit.core.common.mixed_precision import mixed_precisio
 from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig, QuantizationErrorMethod, DEFAULTCONFIG
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
-from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig, MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core.keras.kpi_data_facade import keras_kpi_data
 from model_compression_toolkit.core.pytorch.kpi_data_facade import pytorch_kpi_data

model_compression_toolkit/core/common/framework_implementation.py

@@ -18,7 +18,7 @@ from typing import Callable, Any, List, Tuple, Dict
 import numpy as np
 
 from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS
-from model_compression_toolkit.core import MixedPrecisionQuantizationConfigV2
+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
@@ -304,7 +304,7 @@ class FrameworkImplementation(ABC):
     @abstractmethod
     def get_sensitivity_evaluator(self,
                                   graph: Graph,
-                                  quant_config: MixedPrecisionQuantizationConfigV2,
+                                  quant_config: MixedPrecisionQuantizationConfig,
                                   representative_data_gen: Callable,
                                   fw_info: FrameworkInfo,
                                   hessian_info_service: HessianInfoService = None,

model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py

@@ -13,16 +13,13 @@
 # limitations under the License.
 # ==============================================================================
 
-from enum import Enum
-from typing import List, Callable, Tuple
+from typing import List, Callable
 
-from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.mixed_precision.distance_weighting import get_average_weights
-from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig, DEFAULTCONFIG
 from model_compression_toolkit.core.common.similarity_analyzer import compute_mse
 
 
-class MixedPrecisionQuantizationConfigV2:
+class MixedPrecisionQuantizationConfig:
 
     def __init__(self,
                  compute_distance_fn: Callable = None,
@@ -36,8 +33,6 @@ class MixedPrecisionQuantizationConfigV2:
                  metric_normalization_threshold: float = 1e10):
         """
         Class with mixed precision parameters to quantize the input model.
-        Unlike QuantizationConfig, number of bits for quantization is a list of possible bit widths to
-        support mixed-precision model quantization.
 
         Args:
             compute_distance_fn (Callable): Function to compute a distance between two tensors.
@@ -70,67 +65,6 @@ class MixedPrecisionQuantizationConfigV2:
         self.metric_normalization_threshold = metric_normalization_threshold
 
 
-class MixedPrecisionQuantizationConfig(QuantizationConfig):
-
-    def __init__(self,
-                 qc: QuantizationConfig = DEFAULTCONFIG,
-                 compute_distance_fn: Callable = compute_mse,
-                 distance_weighting_method: Callable = get_average_weights,
-                 num_of_images: int = 32,
-                 configuration_overwrite: List[int] = None,
-                 num_interest_points_factor: float = 1.0):
-        """
-        Class to wrap all different parameters the library quantize the input model according to.
-        Unlike QuantizationConfig, number of bits for quantization is a list of possible bit widths to
-        support mixed-precision model quantization.
-
-        Args:
-            qc (QuantizationConfig): QuantizationConfig object containing parameters of how the model should be quantized.
-            compute_distance_fn (Callable): Function to compute a distance between two tensors.
-            distance_weighting_method (Callable): Function to use when weighting the distances among different layers when computing the sensitivity 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: A multiplication factor between zero and one (represents percentage) to reduce the number of interest points used to calculate the distance metric.
-
-        """
-
-        super().__init__(**qc.__dict__)
-        self.compute_distance_fn = compute_distance_fn
-        self.distance_weighting_method = distance_weighting_method
-        self.num_of_images = num_of_images
-        self.configuration_overwrite = configuration_overwrite
-
-        assert 0.0 < num_interest_points_factor <= 1.0, "num_interest_points_factor should represent a percentage of " \
-                                                        "the base set of interest points that are required to be " \
-                                                        "used for mixed-precision metric evaluation, " \
-                                                        "thus, it should be between 0 to 1"
-        self.num_interest_points_factor = num_interest_points_factor
-
-    def separate_configs(self) -> Tuple[QuantizationConfig, MixedPrecisionQuantizationConfigV2]:
-        """
-        A function to separate the old MixedPrecisionQuantizationConfig into QuantizationConfig
-        and MixedPrecisionQuantizationConfigV2
-
-        Returns: QuantizationConfig, MixedPrecisionQuantizationConfigV2
-
-        """
-        _dummy_quant_config = QuantizationConfig()
-        _dummy_mp_config_experimental = MixedPrecisionQuantizationConfigV2()
-        qc_dict = {}
-        mp_dict = {}
-        for k, v in self.__dict__.items():
-            if hasattr(_dummy_quant_config, k):
-                qc_dict.update({k: v})
-            elif hasattr(_dummy_mp_config_experimental, k):
-                mp_dict.update({k: v})
-            else:
-                Logger.error(f'Attribute "{k}" mismatch: exists in MixedPrecisionQuantizationConfig but not in '
-                             f'MixedPrecisionQuantizationConfigV2')  # pragma: no cover
-
-        return QuantizationConfig(**qc_dict), MixedPrecisionQuantizationConfigV2(**mp_dict)
-
-
 # Default quantization configuration the library use.
-DEFAULT_MIXEDPRECISION_CONFIG = MixedPrecisionQuantizationConfig(DEFAULTCONFIG,
-                                                                 compute_mse,
-                                                                 get_average_weights)
+DEFAULT_MIXEDPRECISION_CONFIG = MixedPrecisionQuantizationConfig(compute_distance_fn=compute_mse,
+                                                                 distance_weighting_method=get_average_weights)

model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py

@@ -18,7 +18,7 @@ from enum import Enum
 import numpy as np
 from typing import List, Callable, Dict
 
-from model_compression_toolkit.core import MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core import MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.core.common.hessian import HessianInfoService
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI, KPITarget
@@ -48,7 +48,7 @@ def search_bit_width(graph_to_search_cfg: Graph,
                      fw_info: FrameworkInfo,
                      fw_impl: FrameworkImplementation,
                      target_kpi: KPI,
-                     mp_config: MixedPrecisionQuantizationConfigV2,
+                     mp_config: MixedPrecisionQuantizationConfig,
                      representative_data_gen: Callable,
                      search_method: BitWidthSearchMethod = BitWidthSearchMethod.INTEGER_PROGRAMMING,
                      hessian_info_service: HessianInfoService=None) -> List[int]:

model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py

@@ -18,7 +18,7 @@ import numpy as np
 from typing import Callable, Any, List, Tuple
 
 from model_compression_toolkit.constants import AXIS, HESSIAN_OUTPUT_ALPHA
-from model_compression_toolkit.core import FrameworkInfo, MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core import FrameworkInfo, MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 
@@ -37,7 +37,7 @@ class SensitivityEvaluation:
 
     def __init__(self,
                  graph: Graph,
-                 quant_config: MixedPrecisionQuantizationConfigV2,
+                 quant_config: MixedPrecisionQuantizationConfig,
                  representative_data_gen: Callable,
                  fw_info: FrameworkInfo,
                  fw_impl: Any,

model_compression_toolkit/core/common/pruning/memory_calculator.py

@@ -307,7 +307,25 @@ class MemoryCalculator:
             total_params += w.size
 
         # Adjust the total parameter count if padded channels are to be included.
-        num_oc = np.sum(output_mask) if output_mask is not None else node.output_shape[-1]
+        if output_mask is not None:
+            num_oc = np.sum(output_mask)
+        else:
+            # Get the node channel axis from framework info
+            channel_axis = self.fw_info.out_channel_axis_mapping.get(node.type)
+            if channel_axis is None:
+                Logger.error("Channel axis is not defined")
+
+            # Check if node.output_shape is a list of lists.
+            # In this case make sure all the out channels are the same value
+            if all(isinstance(sublist, list) for sublist in node.output_shape):
+                compare_value = node.output_shape[0][channel_axis]
+                if all(len(sublist) > channel_axis and sublist[channel_axis] == compare_value for sublist in node.output_shape):
+                    num_oc = compare_value
+                else:
+                    Logger.error("Number of out channels are not the same for all outputs of the node")
+            else:
+                num_oc = node.output_shape[channel_axis]
+
         if include_padded_channels:
             total_params = self.get_node_nparams_with_padded_channels(node, total_params, num_oc, node.get_simd())
 

model_compression_toolkit/core/common/quantization/core_config.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig
 from model_compression_toolkit.core.common.quantization.debug_config import DebugConfig
-from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig
 
 
 class CoreConfig:
@@ -23,14 +23,14 @@ class CoreConfig:
     """
     def __init__(self,
                  quantization_config: QuantizationConfig = QuantizationConfig(),
-                 mixed_precision_config: MixedPrecisionQuantizationConfigV2 = None,
+                 mixed_precision_config: MixedPrecisionQuantizationConfig = None,
                  debug_config: DebugConfig = DebugConfig()
                  ):
         """
 
         Args:
             quantization_config (QuantizationConfig): Config for quantization.
-            mixed_precision_config (MixedPrecisionQuantizationConfigV2): Config for mixed precision quantization (optional, default=None).
+            mixed_precision_config (MixedPrecisionQuantizationConfig): Config for mixed precision quantization (optional, default=None).
             debug_config (DebugConfig): Config for debugging and editing the network quantization process.
         """
         self.quantization_config = quantization_config

model_compression_toolkit/core/common/quantization/quantization_fn_selection.py

@@ -17,7 +17,6 @@ from collections.abc import Callable
 from functools import partial
 
 from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
-from model_compression_toolkit.core.common.quantization.quantizers.kmeans_quantizer import kmeans_quantizer
 from model_compression_toolkit.core.common.quantization.quantizers.lut_kmeans_quantizer import lut_kmeans_quantizer
 from model_compression_toolkit.core.common.quantization.quantizers.uniform_quantizers import power_of_two_quantizer, \
     symmetric_quantizer, uniform_quantizer
@@ -40,8 +39,6 @@ def get_weights_quantization_fn(weights_quantization_method: QuantizationMethod)
         quantizer_fn = symmetric_quantizer
     elif weights_quantization_method == QuantizationMethod.UNIFORM:
         quantizer_fn = uniform_quantizer
-    elif weights_quantization_method == QuantizationMethod.KMEANS:
-        quantizer_fn = kmeans_quantizer
     elif weights_quantization_method in [QuantizationMethod.LUT_POT_QUANTIZER, QuantizationMethod.LUT_SYM_QUANTIZER]:
         quantizer_fn = lut_kmeans_quantizer
     else:

model_compression_toolkit/core/common/quantization/quantization_params_fn_selection.py

@@ -18,7 +18,6 @@ from functools import partial
 
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
-from model_compression_toolkit.core.common.quantization.quantization_params_generation.kmeans_params import kmeans_tensor
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.lut_kmeans_params import \
     lut_kmeans_tensor, lut_kmeans_histogram
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.symmetric_selection import \
@@ -70,8 +69,6 @@ def get_weights_quantization_params_fn(weights_quantization_method: Quantization
         params_fn = symmetric_selection_tensor
     elif weights_quantization_method == QuantizationMethod.UNIFORM:
         params_fn = uniform_selection_tensor
-    elif weights_quantization_method == QuantizationMethod.KMEANS:
-        params_fn = kmeans_tensor
     elif weights_quantization_method == QuantizationMethod.LUT_POT_QUANTIZER:
         params_fn = partial(lut_kmeans_tensor, is_symmetric=False)
     elif weights_quantization_method == QuantizationMethod.LUT_SYM_QUANTIZER:

model_compression_toolkit/core/common/quantization/quantization_params_generation/__init__.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.power_of_two_selection import power_of_two_no_clipping_selection_min_max, \
     power_of_two_selection_histogram, power_of_two_selection_tensor
-from model_compression_toolkit.core.common.quantization.quantization_params_generation.kmeans_params import kmeans_tensor
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.lut_kmeans_params import lut_kmeans_tensor
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.symmetric_selection import symmetric_no_clipping_selection_min_max
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.uniform_selection import uniform_no_clipping_selection_min_max

model_compression_toolkit/core/keras/keras_implementation.py

@@ -57,7 +57,7 @@ else:
         Concatenate, Add
     from keras.layers.core import TFOpLambda
 
-from model_compression_toolkit.core import QuantizationConfig, FrameworkInfo, CoreConfig, MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core import QuantizationConfig, FrameworkInfo, CoreConfig, MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
@@ -355,7 +355,7 @@ class KerasImplementation(FrameworkImplementation):
 
     def get_sensitivity_evaluator(self,
                                   graph: Graph,
-                                  quant_config: MixedPrecisionQuantizationConfigV2,
+                                  quant_config: MixedPrecisionQuantizationConfig,
                                   representative_data_gen: Callable,
                                   fw_info: FrameworkInfo,
                                   disable_activation_for_metric: bool = False,

model_compression_toolkit/core/keras/kpi_data_facade.py

@@ -14,8 +14,7 @@
 # ==============================================================================
 
 from typing import Callable
-
-from model_compression_toolkit.core import CoreConfig, MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core import MixedPrecisionQuantizationConfig, CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import TENSORFLOW
@@ -36,7 +35,7 @@ if FOUND_TF:
 
     def keras_kpi_data(in_model: Model,
                        representative_data_gen: Callable,
-                       core_config: CoreConfig,
+                       core_config: CoreConfig = CoreConfig(mixed_precision_config=MixedPrecisionQuantizationConfig()),
                        fw_info: FrameworkInfo = DEFAULT_KERAS_INFO,
                        target_platform_capabilities: TargetPlatformCapabilities = KERAS_DEFAULT_TPC) -> KPI:
         """
@@ -73,9 +72,9 @@ if FOUND_TF:
 
         """
 
-        if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfigV2):
-            Logger.error("KPI data computation can't be executed without MixedPrecisionQuantizationConfigV2 object."
-                         "Given quant_config is not of type MixedPrecisionQuantizationConfigV2.")
+        if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfig):
+            Logger.error("KPI data computation can't be executed without MixedPrecisionQuantizationConfig object."
+                         "Given quant_config is not of type MixedPrecisionQuantizationConfig.")
 
         fw_impl = KerasImplementation()