mct-nightly 2.2.0.20241025.505__py3-none-any.whl → 2.2.0.20241027.532__py3-none-any.whl

{mct_nightly-2.2.0.20241025.505.dist-info → mct_nightly-2.2.0.20241027.532.dist-info}/METADATA

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

{mct_nightly-2.2.0.20241025.505.dist-info → mct_nightly-2.2.0.20241027.532.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=V06niKRfU-gteePjgcRX-Y6QEsEyAyHzzvzwawNEU3U,1573
+model_compression_toolkit/__init__.py,sha256=e7Hzdcg92ykZG1qmS--aAUq-lsyc1Xp_WN__EKEbiDk,1573
 model_compression_toolkit/constants.py,sha256=i4wYheBkIdQmsQA-axIpcT3YiSO1USNc-jaNiNE8w6E,3920
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
@@ -347,23 +347,24 @@ model_compression_toolkit/gptq/__init__.py,sha256=pEgkJvmf05KSw70iLDTz_6LI_2Oi5L
 model_compression_toolkit/gptq/runner.py,sha256=La12JTYjWyJW0YW4Al4TP1_Xi4JWBCEKw6FR_JQsxe0,5982
 model_compression_toolkit/gptq/common/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/gptq/common/gptq_config.py,sha256=Z6T5B3q4k2Tlr2bBWvC6TAF3d2opyA7ZT_D_mz6D1_0,6297
-model_compression_toolkit/gptq/common/gptq_constants.py,sha256=QSm6laLkIV0LYmU0BLtmKp3Fi3SqDfbncFQWOGA1cGU,611
+model_compression_toolkit/gptq/common/gptq_constants.py,sha256=D1x2n4-NdAx6g_1Wc2hwwh4vX9vmx5VnQWN26H107kg,766
 model_compression_toolkit/gptq/common/gptq_framework_implementation.py,sha256=n3mSf4J92kFjekzyGyrJULylI-8Jf5OVWJ5AFoVnEx0,1266
 model_compression_toolkit/gptq/common/gptq_graph.py,sha256=-bL5HhPcKqV8nj4dZPXc5QmQJbFBel6etrioikP0tEo,3039
 model_compression_toolkit/gptq/common/gptq_training.py,sha256=tt4O8PjSChquzl4c6NojvQWZmvCdTxcMLtmEVIGx1ns,13252
+model_compression_toolkit/gptq/common/gradual_activation_quantization.py,sha256=EgpzMs_aDoB0wQiTagqvcxCTfrgNUuCfdXEXmfNiyb0,3780
+model_compression_toolkit/gptq/common/regularization_factory.py,sha256=hyunpXepVeHyoAFJw6zNLK-3ZHBmiut3lmNisJN_L3E,2514
 model_compression_toolkit/gptq/keras/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/gptq/keras/gptq_keras_implementation.py,sha256=axBwnCSjq5xk-xGymOwSOqjp39It-CVtGcCTRTf0E_4,1248
 model_compression_toolkit/gptq/keras/gptq_loss.py,sha256=rbRkF15MYd6nq4G49kcjb_dPTa-XNq9cTkrb93mXawo,6241
-model_compression_toolkit/gptq/keras/gptq_training.py,sha256=tFHucF7YHKtHmYGkdMpqSf14H9c7x60Il7ZTMNXSesE,19751
+model_compression_toolkit/gptq/keras/gptq_training.py,sha256=TEWqAU8JZnZVZ-dIkINA0x1NmSrYpEkXTdG835JdKnI,20848
 model_compression_toolkit/gptq/keras/graph_info.py,sha256=MKIfrRTRH3zCuxCR1g9ZVIFyuSSr0e0sDybqh4LDM7E,4672
-model_compression_toolkit/gptq/keras/quantization_facade.py,sha256=iSHnMEdoIqHYqLCTsdK8uxhKbZuuaDOu_BeQ10Z492U,15715
+model_compression_toolkit/gptq/keras/quantization_facade.py,sha256=DhEEpW0rK4JRdk5WQlN-_DOUuzlwOBqpiwTBOySjn2g,16820
 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
-model_compression_toolkit/gptq/keras/quantizer/quantization_builder.py,sha256=FmK5cPwgLAzrDjHTWf_vbRO5s70S7iwpnjnlqEQTuGE,4408
-model_compression_toolkit/gptq/keras/quantizer/regularization_factory.py,sha256=guf7ygnLsZeWnTDz4yJdE2iTkd1oE0uQAZwKnGV3OAk,1957
+model_compression_toolkit/gptq/keras/quantizer/quantization_builder.py,sha256=rst-u5EB9Xss4ndKqi297WvZ-9RVee2TAUVFelPVKhU,4663
 model_compression_toolkit/gptq/keras/quantizer/soft_rounding/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
-model_compression_toolkit/gptq/keras/quantizer/soft_rounding/soft_quantizer_reg.py,sha256=qUuMKysUpjWYjNbchFuyb_UFwzV1HL7R3Y7o0Z5rf60,4016
+model_compression_toolkit/gptq/keras/quantizer/soft_rounding/soft_quantizer_reg.py,sha256=REO-pIXpT4ZuJzhizvQjz6vn7Vxnq7k0KvikuQ4FDkE,2769
 model_compression_toolkit/gptq/keras/quantizer/soft_rounding/symmetric_soft_quantizer.py,sha256=BBSDWLmeywjSM5N6oJkMgcuo7zrXTesB4zLwRGG8QB0,12159
 model_compression_toolkit/gptq/keras/quantizer/soft_rounding/uniform_soft_quantizer.py,sha256=pyhlVpoauHM-zuixHsIGPHFgQoXppL8TlDFCjPE2RuY,10377
 model_compression_toolkit/gptq/keras/quantizer/ste_rounding/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
@@ -371,15 +372,13 @@ model_compression_toolkit/gptq/keras/quantizer/ste_rounding/symmetric_ste.py,sha
 model_compression_toolkit/gptq/pytorch/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/gptq/pytorch/gptq_loss.py,sha256=_07Zx_43bnNokwR5S8phIqeu5-_7_5VBT4DT-FCw7Do,3892
 model_compression_toolkit/gptq/pytorch/gptq_pytorch_implementation.py,sha256=tECPTavxn8EEwgLaP2zvxdJH6Vg9jC0YOIMJ7857Sdc,1268
-model_compression_toolkit/gptq/pytorch/gptq_training.py,sha256=QBxTnwVvLyZDTdpkR81wjj9o5aGtmp9qiBt5FR8ImJ0,21777
+model_compression_toolkit/gptq/pytorch/gptq_training.py,sha256=2KwJFlJj6hFJClsJbC9aaWDAGbZUNDbSx1d-QX4LShc,22132
 model_compression_toolkit/gptq/pytorch/graph_info.py,sha256=4mVM-VvnBaA64ACVdOe6wTGHdMSa2UTLIUe7nACLcdo,4008
-model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=yv2DWPWpFVRmtB_FhcRwnLUumyPPHC_hHaMxeQBTQ1k,16333
+model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=lY7_lNtS1SqaaJ0gc6C7_HO71bBalsxQY37QQlWpu70,15479
 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/gradual_activation_quantization.py,sha256=nngu2TeXjngkqt_6-wciFmCvo-dbpeh_tJJxBV_cfHk,3686
 model_compression_toolkit/gptq/pytorch/quantizer/quant_utils.py,sha256=OocYYRqvl7rZ37QT0hTzfJnWGiNCPskg7cziTlR7TRk,3893
 model_compression_toolkit/gptq/pytorch/quantizer/quantization_builder.py,sha256=5EyAzvlU01vLyXmMwY_8dNyb7GwYktXmnrvUON8n8WI,4696
-model_compression_toolkit/gptq/pytorch/quantizer/regularization_factory.py,sha256=H6pARLK-jq3cKoaipY0SK9wMGrqy6CSEZTk14KdrKA0,2105
 model_compression_toolkit/gptq/pytorch/quantizer/soft_rounding/__init__.py,sha256=lNJ29DYxaLUPDstRDA1PGI5r9Fulq_hvrZMlhst1Z5g,697
 model_compression_toolkit/gptq/pytorch/quantizer/soft_rounding/soft_quantizer_reg.py,sha256=f7B95Bx-MX-HKheqAUn1GG8cVHFI2ldFReXrUPwk2tY,3002
 model_compression_toolkit/gptq/pytorch/quantizer/soft_rounding/symmetric_soft_quantizer.py,sha256=kLVQC1hXzDpP4Jx7AwnA764oGnY5AMEuvUUhAvhz09M,12347
@@ -407,7 +406,7 @@ model_compression_toolkit/qat/keras/quantizer/base_keras_qat_weight_quantizer.py
 model_compression_toolkit/qat/keras/quantizer/quant_utils.py,sha256=cBULOgWUodcBO1lHevZggdTevuDYI6tQceV86U2x6DA,2543
 model_compression_toolkit/qat/keras/quantizer/quantization_builder.py,sha256=hGizGBbOGZpD-w3wg-LlehUYJDWLk91VUdfVwwG2Z78,5882
 model_compression_toolkit/qat/keras/quantizer/lsq/__init__.py,sha256=lNJ29DYxaLUPDstRDA1PGI5r9Fulq_hvrZMlhst1Z5g,697
-model_compression_toolkit/qat/keras/quantizer/lsq/symmetric_lsq.py,sha256=dtkS0mpjvJntAxpOi-BJx-pCeBF2ReKKeH7y2uwzpH0,6756
+model_compression_toolkit/qat/keras/quantizer/lsq/symmetric_lsq.py,sha256=ujxb7hYHj25QRwu7SScP8BZXWQdh61knzk68KRIdOp0,6501
 model_compression_toolkit/qat/keras/quantizer/lsq/uniform_lsq.py,sha256=vGUs9b0IHTydCA5tN7iekuhf1LHNgIrSF5sXMD1WsSI,6476
 model_compression_toolkit/qat/keras/quantizer/ste_rounding/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/qat/keras/quantizer/ste_rounding/symmetric_ste.py,sha256=pFkrjtlavCniswcO3-Djlh6a_Hz1rrcEa7Z5wTGVRCU,8270
@@ -493,6 +492,7 @@ model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/v1/
 model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/v1/tpc_pytorch.py,sha256=YVJJvqGPBdkKnug99p9bjqtbfecDXZKIB2iWVCe7RUY,5960
 model_compression_toolkit/trainable_infrastructure/__init__.py,sha256=uewpvlPkH9mBFt8IxoAgIfz6iEcvWbOImm_fb6_BxD8,1543
 model_compression_toolkit/trainable_infrastructure/common/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
+model_compression_toolkit/trainable_infrastructure/common/annealing_schedulers.py,sha256=qm2_wa61nga08Jdcl3RkgTsJ0zyHNjZ_A6I2--oVOig,2455
 model_compression_toolkit/trainable_infrastructure/common/base_trainable_quantizer.py,sha256=IF50ASBUvVrOVqlJ1nHNxZxKXSuCanjhUX0YjMB-rRg,7946
 model_compression_toolkit/trainable_infrastructure/common/constants.py,sha256=HN120boJxAnEXNrLSj-o_s-VX4o6C-1ap_KZ4840sd0,875
 model_compression_toolkit/trainable_infrastructure/common/get_quantizer_config.py,sha256=Jxd4IjS_t0FwnA_S_WmZeVbh4VM6Da9ahKGPLp6ZhQo,6983
@@ -500,7 +500,9 @@ model_compression_toolkit/trainable_infrastructure/common/get_quantizers.py,sha2
 model_compression_toolkit/trainable_infrastructure/common/quant_utils.py,sha256=zdiew1jwR7tUKm9XWlHnAPxIZsAdKqbzzC2vH02j5wA,1505
 model_compression_toolkit/trainable_infrastructure/common/trainable_quantizer_config.py,sha256=My5Wz34jPOyh8z33OTpKnOobRB0cpO_Qgmtsd5lizHo,4791
 model_compression_toolkit/trainable_infrastructure/common/training_method.py,sha256=LUoeJkloowhZKuHTiOfzjmSUn2G-4of11-rbnL-h0P4,1194
+model_compression_toolkit/trainable_infrastructure/common/util.py,sha256=oKuWi7E07a8zv5x9auhBugYE2RUQ7ojDh2XCs5koYJY,1090
 model_compression_toolkit/trainable_infrastructure/keras/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
+model_compression_toolkit/trainable_infrastructure/keras/annealing_schedulers.py,sha256=sISNVxPsdm-Nd95PhoPSJ-2tFpINGlfrU7ZXaCByI-o,1278
 model_compression_toolkit/trainable_infrastructure/keras/base_keras_quantizer.py,sha256=LBc26z8pkpbcdKMTxpNBg5IyChLreHQ1lRgCVjNE37o,4202
 model_compression_toolkit/trainable_infrastructure/keras/config_serialization.py,sha256=txdWXdZoHazg-3MDPb9P-oXRM92LRn2G_8woEplwKaI,4360
 model_compression_toolkit/trainable_infrastructure/keras/load_model.py,sha256=DJHibcLo-UCuHV6UPLeVd7dKmPfkGXEiLqCCqvQrISM,3769
@@ -515,10 +517,9 @@ model_compression_toolkit/trainable_infrastructure/keras/activation_quantizers/s
 model_compression_toolkit/trainable_infrastructure/keras/activation_quantizers/ste/symmetric_ste.py,sha256=THY5eZ_69D1yzkXLhLg84ON_deNUAD_qMJ6A5C5znDM,7359
 model_compression_toolkit/trainable_infrastructure/keras/activation_quantizers/ste/uniform_ste.py,sha256=XEypYorBnSBLj6sh1pHCNaSjeCToYVlERWIHxUoXvuc,5733
 model_compression_toolkit/trainable_infrastructure/pytorch/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
-model_compression_toolkit/trainable_infrastructure/pytorch/annealing_schedulers.py,sha256=IdUBpZUcOXHLPp2OhwbO_Kytee3OTVuy2032N-tm694,1686
+model_compression_toolkit/trainable_infrastructure/pytorch/annealing_schedulers.py,sha256=W5NPQiwIAd2dpaoU9WfRwSt0ljrrePj5lwPk6d1yVwQ,1333
 model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py,sha256=lWc5EG3ptrP85n69EHGKFkIadnrKEBMKnB5YXQ5AmXo,2745
 model_compression_toolkit/trainable_infrastructure/pytorch/quantizer_utils.py,sha256=1yOXKghUYfw2hmzbqTuNagIXBoM-wR2bP-ul66-mnDw,7767
-model_compression_toolkit/trainable_infrastructure/pytorch/util.py,sha256=oKuWi7E07a8zv5x9auhBugYE2RUQ7ojDh2XCs5koYJY,1090
 model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/__init__.py,sha256=73CXhqqNTvDpsvlJXclrGJq-vsCUYCI64ILu1y2mtvw,1056
 model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/base_activation_quantizer.py,sha256=X6E6mewWQot_aAkz3UxW5X0-Fjl_aMMjs3A-Af5eL6w,972
 model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/lsq/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
@@ -557,16 +558,20 @@ tests_pytest/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
 tests_pytest/keras/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
 tests_pytest/keras/core/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
 tests_pytest/keras/core/test_data_util.py,sha256=XSoPu_ci1xy2EtK-3OWGpESr-Meg1GDaxuSvcj3yt-w,3915
+tests_pytest/keras/gptq/__init__.py,sha256=pKAdbTCFM_2BrZXUtTIw0ouKotrWwUDF_hP3rPwCM2k,696
+tests_pytest/keras/gptq/test_gradual_act_quantization.py,sha256=iwKaLI7QQ8H3qj6zmwwfd2ZOwRcCr8T-v_4llSh_chM,4804
+tests_pytest/keras/trainable_infrastructure/__init__.py,sha256=DvaMXJtJZHAqOm96NdfBiNQsbN2sc9bG2kkyY-mpPh8,710
+tests_pytest/keras/trainable_infrastructure/test_linear_annealing.py,sha256=dZjrMHVIiEVRNDYR3a4lZaXF2ElxFx32KAXXQvDz-v8,1793
 tests_pytest/pytorch/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
 tests_pytest/pytorch/core/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
 tests_pytest/pytorch/core/test_data_util.py,sha256=Bg3c21YVfXE1SAUlTao553gXcITTKF4CPeKtl3peBTE,5604
 tests_pytest/pytorch/gptq/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
 tests_pytest/pytorch/gptq/test_annealing_cfg.py,sha256=hGC7L6mp3N1ygcJ3OctgS_Fz2JY75q5aswolJkbHkZM,2208
-tests_pytest/pytorch/gptq/test_gradual_act_quantization.py,sha256=tI01aFIUaiCILL5Qn--p1E_rLBUelxLdSY3k52lwcx0,4594
+tests_pytest/pytorch/gptq/test_gradual_act_quantization.py,sha256=Dg2cg1X8u9Jxm7Y6tlZIGH81EPoW_vYorcdDExdj02w,4630
 tests_pytest/pytorch/trainable_infrastructure/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
-tests_pytest/pytorch/trainable_infrastructure/test_linear_annealing.py,sha256=eNOpSp0GoLxtEdiRypBp8jaujXfdNxBwKh5Rd-P7WLs,1786
-mct_nightly-2.2.0.20241025.505.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.2.0.20241025.505.dist-info/METADATA,sha256=Sp107NNGeo7gOmqCwchsuQSnMQ5UR5VkkioWUUCgsak,20830
-mct_nightly-2.2.0.20241025.505.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-mct_nightly-2.2.0.20241025.505.dist-info/top_level.txt,sha256=csdfSXhtRnpWYRzjZ-dRLIhOmM2TEdVXUxG05A5fgb8,39
-mct_nightly-2.2.0.20241025.505.dist-info/RECORD,,
+tests_pytest/pytorch/trainable_infrastructure/test_linear_annealing.py,sha256=zErt9tOu7oupjpv08cvd1Cxvdk9qvP7GMUP6EhefK0c,1814
+mct_nightly-2.2.0.20241027.532.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.2.0.20241027.532.dist-info/METADATA,sha256=F0SYwH78ncR2_2Cx8EnUV0f_eydBxIdBjGnOd22oxqM,20830
+mct_nightly-2.2.0.20241027.532.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+mct_nightly-2.2.0.20241027.532.dist-info/top_level.txt,sha256=csdfSXhtRnpWYRzjZ-dRLIhOmM2TEdVXUxG05A5fgb8,39
+mct_nightly-2.2.0.20241027.532.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.20241025.000505"
+__version__ = "2.2.0.20241027.000532"

model_compression_toolkit/gptq/common/gptq_constants.py

@@ -22,4 +22,11 @@ SOFT_ROUNDING_ZETA = 1.1
 
 # GPTQ config constant
 QUANT_PARAM_LEARNING_STR = 'quantization_parameter_learning'
-MAX_LSB_STR = 'max_lsbs_change_map'
+MAX_LSB_STR = 'max_lsbs_change_map'
+
+# GPTQ learning hyperparameters
+LR_DEFAULT = 3e-2
+LR_REST_DEFAULT = 1e-4
+LR_BIAS_DEFAULT = 1e-3
+LR_QUANTIZATION_PARAM_DEFAULT = 1e-3
+GPTQ_MOMENTUM = 0.9

model_compression_toolkit/gptq/{pytorch/quantizer → common}/gradual_activation_quantization.py

@@ -13,23 +13,23 @@
 # limitations under the License.
 # ==============================================================================
 from functools import partial
-from typing import Callable
+from typing import Callable, Any
 
 from model_compression_toolkit.gptq import GradientPTQConfig, QFractionLinearAnnealingConfig
-from model_compression_toolkit.trainable_infrastructure import BasePytorchTrainableQuantizer
-
-from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import LinearAnnealingScheduler
+from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import BaseTrainableQuantizer
 
 
 def get_gradual_activation_quantizer_wrapper_factory(gptq_config: GradientPTQConfig,
-                                                     get_total_grad_steps_fn: Callable[[], int]) \
-        -> Callable[[BasePytorchTrainableQuantizer], 'GradualActivationQuantizerWrapper']:
+                                                     get_total_grad_steps_fn: Callable[[], int],
+                                                     fw_linear_annealing_scheduler: type) \
+        -> Callable[[Any], 'GradualActivationQuantizerWrapper']:
     """
     Get a factory for 'GradualActivationQuantizerWrapper'.
 
     Args:
         gptq_config: GPTQ configuration.
         get_total_grad_steps_fn: a callable to obtain the total expected number of gradient steps.
+        fw_linear_annealing_scheduler: LinearAnnealingScheduler implementation of the framework (tf/pytorch).
 
     Returns:
         A factory function to build 'GradualActivationQuantizerWrapper' from Quantizer.
@@ -40,9 +40,9 @@ def get_gradual_activation_quantizer_wrapper_factory(gptq_config: GradientPTQCon
     annealing_cfg = gptq_config.gradual_activation_quantization_config.q_fraction_scheduler_policy
     if isinstance(annealing_cfg, QFractionLinearAnnealingConfig):
         t_end = annealing_cfg.end_step or get_total_grad_steps_fn()
-        factor_scheduler = LinearAnnealingScheduler(t_start=annealing_cfg.start_step, t_end=t_end,
-                                                    initial_val=annealing_cfg.initial_q_fraction,
-                                                    target_val=annealing_cfg.target_q_fraction)
+        factor_scheduler = fw_linear_annealing_scheduler(t_start=annealing_cfg.start_step, t_end=t_end,
+                                                         initial_val=annealing_cfg.initial_q_fraction,
+                                                         target_val=annealing_cfg.target_q_fraction)
     else:
         raise ValueError(f'Unknown annealing policy {annealing_cfg}')
 
@@ -64,7 +64,7 @@ class GradualActivationQuantizerWrapper:
         quantizer: quantizer to wrap.
         q_fraction_scheduler: a callable that accepts a gradient step and returns the corresponding quantized fraction.
     """
-    def __init__(self, quantizer: BasePytorchTrainableQuantizer, q_fraction_scheduler: Callable[[int], float]):
+    def __init__(self, quantizer: BaseTrainableQuantizer, q_fraction_scheduler: Callable[[int], float]):
         self.quantizer = quantizer
         self.q_fraction_scheduler = q_fraction_scheduler
         self.step_cnt = 0

model_compression_toolkit/gptq/{pytorch/quantizer → common}/regularization_factory.py

@@ -12,17 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Callable
 
-from model_compression_toolkit.gptq import RoundingType, GradientPTQConfig
-from model_compression_toolkit.gptq.pytorch.quantizer.soft_rounding.soft_quantizer_reg import \
-    SoftQuantizerRegularization
-from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import LinearAnnealingScheduler
+from tqdm import tqdm
+from typing import Callable, Type
 
+from model_compression_toolkit.gptq import RoundingType, GradientPTQConfig
 
+# Common warmup fraction
 WARMUP_STEP_FRACTION = 0.2
 
-def get_regularization(gptq_config: GradientPTQConfig, get_total_grad_steps_fn: Callable[[], int]) -> Callable:
+
+def get_regularization(gptq_config: GradientPTQConfig,
+                       get_total_grad_steps_fn: Callable[[], int],
+                       SoftQuantizerRegularizationFWClass: Type,
+                       LinearAnnealingSchedulerFWClass: Type) -> Callable:
     """
     Returns a function that computes the regularization term for GPTQ training based on the given
     rounding type in the GPTQ configuration.
@@ -30,15 +33,26 @@ def get_regularization(gptq_config: GradientPTQConfig, get_total_grad_steps_fn:
     Args:
         gptq_config: A GPTQ configuration.
         get_total_grad_steps_fn: a callable to obtain the total expected number of gradient steps.
+        SoftQuantizerRegularizationFWClass: The class to use for soft quantizer regularization (framework-specific).
+        LinearAnnealingSchedulerFWClass: The class to use for the annealing scheduler (framework-specific).
 
-    Returns: A function for computing the regularization. If there is no regularization function defined for the given
-        rounding type, then it returns a function that just returns 0.
-
+    Returns:
+        Callable: A function for computing the regularization. If there is no regularization function
+        defined for the given rounding type, then it returns a function that just returns 0.
     """
     if gptq_config.rounding_type == RoundingType.SoftQuantizer:
         total_gradient_steps = get_total_grad_steps_fn()
         t_start = int(WARMUP_STEP_FRACTION * total_gradient_steps)
-        scheduler = LinearAnnealingScheduler(t_start=t_start, t_end=total_gradient_steps, initial_val=20, target_val=2)
-        return SoftQuantizerRegularization(scheduler)
+
+        # Directly initializing the scheduler within the method
+        scheduler = LinearAnnealingSchedulerFWClass(
+            t_start=t_start,
+            t_end=total_gradient_steps,
+            initial_val=20,
+            target_val=2
+        )
+
+        # Return the framework-specific soft quantizer regularization
+        return SoftQuantizerRegularizationFWClass(scheduler)
     else:
         return lambda *args, **kwargs: 0

model_compression_toolkit/gptq/keras/gptq_training.py

@@ -26,9 +26,14 @@ from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder
 from model_compression_toolkit.core.keras.data_util import data_gen_to_dataloader
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
+from model_compression_toolkit.gptq.common.gradual_activation_quantization import \
+    get_gradual_activation_quantizer_wrapper_factory
+from model_compression_toolkit.gptq.common.regularization_factory import get_regularization
 from model_compression_toolkit.gptq.keras.quantizer.quantization_builder import quantization_builder
 from model_compression_toolkit.logger import Logger
 from mct_quantizers import KerasActivationQuantizationHolder
+from model_compression_toolkit.trainable_infrastructure.common.util import get_total_grad_steps
+from model_compression_toolkit.trainable_infrastructure.keras.annealing_schedulers import KerasLinearAnnealingScheduler
 
 if version.parse(tf.__version__) >= version.parse("2.13"):
     from keras.src.engine.base_layer import TensorFlowOpLayer
@@ -41,13 +46,12 @@ from model_compression_toolkit.gptq.common.gptq_training import GPTQTrainer
 from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.gptq.keras.graph_info import get_weights_for_loss, get_gptq_trainable_parameters
-from model_compression_toolkit.gptq.keras.quantizer.regularization_factory import get_regularization
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 import numpy as np
 import copy
 from model_compression_toolkit.core.keras.constants import BIAS, USE_BIAS
-
+from model_compression_toolkit.gptq.keras.quantizer.soft_rounding.soft_quantizer_reg import SoftQuantizerRegularization
 
 class KerasGPTQTrainer(GPTQTrainer):
     """
@@ -78,6 +82,15 @@ class KerasGPTQTrainer(GPTQTrainer):
             hessian_info_service: HessianScoresService for fetching and computing Hessian's approximation scores.
 
         """
+
+        def _get_total_grad_steps():
+            return get_total_grad_steps(representative_data_gen) * gptq_config.n_epochs
+
+        # This must be set before the model building (as it is required for activation holder construction),
+        # which occurs in the base constructor.
+        self.gradual_act_quantizer_wrapper_factory = get_gradual_activation_quantizer_wrapper_factory(
+            gptq_config, _get_total_grad_steps, KerasLinearAnnealingScheduler)
+
         super().__init__(graph_float,
                          graph_quant,
                          gptq_config,
@@ -119,7 +132,10 @@ class KerasGPTQTrainer(GPTQTrainer):
 
         self.weights_for_average_loss = self._get_compare_points_loss_weights()
 
-        self.reg_func = get_regularization(self.gptq_config, representative_data_gen)
+        self.reg_func = get_regularization(self.gptq_config,
+                                           _get_total_grad_steps,
+                                           SoftQuantizerRegularization,
+                                           KerasLinearAnnealingScheduler)
 
     def _get_compare_points_loss_weights(self):
         """ Get compare points weights for the distillation loss. """
@@ -185,14 +201,13 @@ class KerasGPTQTrainer(GPTQTrainer):
         _, activation_quantizers = quantization_builder(n, self.gptq_config) # TODO: split quantizers building into two functions: for weights and activations
 
         # Holder by definition uses a single quantizer for the activation quantization
-        # thus we make sure this is the only possible case (unless it's a node with no activation
-        # quantization, which in this case has an empty list).
-        if len(activation_quantizers) == 1:
-            return KerasActivationQuantizationHolder(activation_quantizers[0])
-
-        Logger.critical(f"'KerasActivationQuantizationHolder' is designed to support a single quantizer, "
-                        f"but {len(activation_quantizers)} quantizers were found for node '{n}'. "
-                        f"Ensure only one quantizer is configured for each node's activation.")
+        # thus we make sure this is the only possible case.
+        if len(activation_quantizers) != 1:
+            Logger.critical(f"'KerasActivationQuantizationHolder' is designed to support a single quantizer, "
+                            f"but {len(activation_quantizers)} quantizers were found for node '{n}'. "
+                            f"Ensure only one quantizer is configured for each node's activation.")
+        quantizer = self.gradual_act_quantizer_wrapper_factory(activation_quantizers[0])
+        return KerasActivationQuantizationHolder(quantizer)
 
     def build_gptq_model(self) -> Tuple[Model, UserInformation]:
         """

model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -14,17 +14,18 @@
 # ==============================================================================
 import copy
 
-from typing import Callable, Tuple
+from typing import Callable, Tuple, Union
 from packaging import version
 
-from model_compression_toolkit.core.common.quantization.quantize_graph_weights import quantize_graph_weights
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
-from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT
+from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT, LR_DEFAULT, LR_REST_DEFAULT, \
+    LR_BIAS_DEFAULT, GPTQ_MOMENTUM
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import TENSORFLOW, ACT_HESSIAN_DEFAULT_BATCH_SIZE
 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
+from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig, GPTQHessianScoresConfig, \
+    GradualActivationQuantizationConfig
 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 MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core import CoreConfig
@@ -32,13 +33,8 @@ from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.gptq.runner import gptq_runner
 from model_compression_toolkit.core.analyzer import analyzer_model_quantization
 from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import TargetPlatformCapabilities
-from model_compression_toolkit.metadata import get_versions_dict, create_model_metadata
+from model_compression_toolkit.metadata import create_model_metadata
 
-LR_DEFAULT = 0.15
-LR_REST_DEFAULT = 1e-4
-LR_BIAS_DEFAULT = 1e-4
-LR_QUANTIZATION_PARAM_DEFAULT = 1e-3
-GPTQ_MOMENTUM = 0.9
 
 if FOUND_TF:
     import tensorflow as tf
@@ -54,25 +50,25 @@ if FOUND_TF:
 
     # As from TF2.9 optimizers package is changed
     if version.parse(tf.__version__) < version.parse("2.9"):
-        from keras.optimizer_v2.optimizer_v2 import OptimizerV2
+        from keras.optimizer_v2.optimizer_v2 import OptimizerV2  # pragma: no cover
     elif version.parse(tf.__version__) < version.parse("2.12"):
-        from keras.optimizers.optimizer_v2.optimizer_v2 import OptimizerV2
+        from keras.optimizers.optimizer_v2.optimizer_v2 import OptimizerV2  # pragma: no cover
     else:
         from tensorflow.python.keras.optimizer_v2.optimizer_v2 import OptimizerV2
 
     DEFAULT_KERAS_TPC = get_target_platform_capabilities(TENSORFLOW, DEFAULT_TP_MODEL)
 
-
     def get_keras_gptq_config(n_epochs: int,
-                              optimizer: OptimizerV2 = tf.keras.optimizers.Adam(learning_rate=LR_DEFAULT),
-                              optimizer_rest: OptimizerV2 = tf.keras.optimizers.Adam(learning_rate=LR_REST_DEFAULT),
+                              optimizer: OptimizerV2 = None,
+                              optimizer_rest: OptimizerV2 = None,
                               loss: Callable = GPTQMultipleTensorsLoss(),
                               log_function: Callable = None,
                               use_hessian_based_weights: bool = True,
                               regularization_factor: float = REG_DEFAULT,
-                              hessian_batch_size: int = ACT_HESSIAN_DEFAULT_BATCH_SIZE) -> GradientPTQConfig:
+                              hessian_batch_size: int = ACT_HESSIAN_DEFAULT_BATCH_SIZE,
+                              gradual_activation_quantization: Union[bool, GradualActivationQuantizationConfig] = False) -> GradientPTQConfig:
         """
-        Create a GradientPTQConfigV2 instance for Keras models.
+        Create a GradientPTQConfig instance for Keras models.
 
         args:
             n_epochs (int): Number of epochs for running the representative dataset for fine-tuning.
@@ -83,9 +79,10 @@ if FOUND_TF:
             use_hessian_based_weights (bool): Whether to use Hessian-based weights for weighted average loss.
             regularization_factor (float): A floating point number that defines the regularization factor.
             hessian_batch_size (int): Batch size for Hessian computation in Hessian-based weights GPTQ.
+            gradual_activation_quantization (bool, GradualActivationQuantizationConfig): If False, GradualActivationQuantization is disabled. If True, GradualActivationQuantization is enabled with the default settings. GradualActivationQuantizationConfig object can be passed to use non-default settings.
 
         returns:
-            a GradientPTQConfigV2 object to use when fine-tuning the quantized model using gptq.
+            a GradientPTQConfig object to use when fine-tuning the quantized model using gptq.
 
         Examples:
 
@@ -94,7 +91,7 @@ if FOUND_TF:
             >>> import model_compression_toolkit as mct
             >>> import tensorflow as tf
 
-            Create a GradientPTQConfigV2 to run for 5 epochs:
+            Create a GradientPTQConfig to run for 5 epochs:
 
             >>> gptq_conf = mct.gptq.get_keras_gptq_config(n_epochs=5)
 
@@ -102,11 +99,24 @@ if FOUND_TF:
 
             >>> gptq_conf = mct.gptq.get_keras_gptq_config(n_epochs=3, optimizer=tf.keras.optimizers.Nadam())
 
-            The configuration can be passed to :func:`~model_compression_toolkit.keras_post_training_quantization` in order to quantize a keras model using gptq.
+            The configuration can be passed to :func:`~model_compression_toolkit.keras_gradient_post_training_quantization` in order to quantize a keras model using gptq.
+
 
         """
+        optimizer = optimizer or tf.keras.optimizers.Adam(learning_rate=LR_DEFAULT)
+        optimizer_rest = optimizer_rest or tf.keras.optimizers.Adam(learning_rate=LR_REST_DEFAULT)
+
         bias_optimizer = tf.keras.optimizers.SGD(learning_rate=LR_BIAS_DEFAULT,
                                                  momentum=GPTQ_MOMENTUM)
+
+        if isinstance(gradual_activation_quantization, bool):
+            gradual_quant_config = GradualActivationQuantizationConfig() if gradual_activation_quantization else None
+        elif isinstance(gradual_activation_quantization, GradualActivationQuantizationConfig):
+            gradual_quant_config = gradual_activation_quantization
+        else:
+            raise TypeError(f'gradual_activation_quantization argument should be bool or '
+                            f'GradualActivationQuantizationConfig, received {type(gradual_activation_quantization)}')
+
         return GradientPTQConfig(n_epochs,
                                  optimizer,
                                  optimizer_rest=optimizer_rest,
@@ -116,7 +126,8 @@ if FOUND_TF:
                                  optimizer_bias=bias_optimizer,
                                  use_hessian_based_weights=use_hessian_based_weights,
                                  regularization_factor=regularization_factor,
-                                 hessian_weights_config=GPTQHessianScoresConfig(hessian_batch_size=hessian_batch_size))
+                                 hessian_weights_config=GPTQHessianScoresConfig(hessian_batch_size=hessian_batch_size),
+                                 gradual_activation_quantization_config=gradual_quant_config)
 
 
     def keras_gradient_post_training_quantization(in_model: Model, representative_data_gen: Callable,
@@ -251,13 +262,13 @@ if FOUND_TF:
 else:
     # If tensorflow is not installed,
     # we raise an exception when trying to use these functions.
-    def get_keras_gptq_config(*args, **kwargs):
+    def get_keras_gptq_config(*args, **kwargs):  # pragma: no cover
         Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
                         "get_keras_gptq_config. The 'tensorflow' package is missing or is "
                         "installed with a version higher than 2.15.")  # pragma: no cover
 
 
-    def keras_gradient_post_training_quantization(*args, **kwargs):
+    def keras_gradient_post_training_quantization(*args, **kwargs):  # pragma: no cover
         Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
                         "keras_gradient_post_training_quantization. The 'tensorflow' package is missing or is "
-                        "installed with a version higher than 2.15.")  # pragma: no cover
+                        "installed with a version higher than 2.15.")

model_compression_toolkit/gptq/keras/quantizer/quantization_builder.py

@@ -16,18 +16,18 @@ from typing import Dict, List, Tuple
 
 from model_compression_toolkit.gptq import GradientPTQConfig
 from model_compression_toolkit.core import common
-from model_compression_toolkit.exporter.model_wrapper.keras.builder.node_to_quantizer import \
-    get_inferable_quantizer_kwargs
 from model_compression_toolkit.gptq.keras.quantizer.base_keras_gptq_quantizer import BaseKerasGPTQTrainableQuantizer
 from mct_quantizers import QuantizationTarget
-from mct_quantizers.common.get_quantizers import get_inferable_quantizer_class
 from mct_quantizers.keras.quantizers import BaseKerasInferableQuantizer
 
 from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.trainable_infrastructure import TrainingMethod
 from model_compression_toolkit.trainable_infrastructure.common.get_quantizer_config import \
-    get_trainable_quantizer_weights_config
+    get_trainable_quantizer_weights_config, get_trainable_quantizer_activation_config
 from model_compression_toolkit.trainable_infrastructure.common.get_quantizers import \
     get_trainable_quantizer_class
+from model_compression_toolkit.trainable_infrastructure.keras.activation_quantizers.base_activation_quantizer import \
+    BaseKerasActivationTrainableQuantizer
 
 
 def quantization_builder(n: common.BaseNode,
@@ -70,12 +70,13 @@ def quantization_builder(n: common.BaseNode,
 
         quant_method = n.final_activation_quantization_cfg.activation_quantization_method
 
-        quantizer_class = get_inferable_quantizer_class(quant_target=QuantizationTarget.Activation,
+        quantizer_class = get_trainable_quantizer_class(quant_target=QuantizationTarget.Activation,
+                                                        quantizer_id=TrainingMethod.STE,
                                                         quant_method=quant_method,
-                                                        quantizer_base_class=BaseKerasInferableQuantizer)
+                                                        quantizer_base_class=BaseKerasActivationTrainableQuantizer)
+        cfg = get_trainable_quantizer_activation_config(n, None)
 
-        kwargs = get_inferable_quantizer_kwargs(n.final_activation_quantization_cfg, QuantizationTarget.Activation)
-
-        activation_quantizers.append(quantizer_class(**kwargs))
+        # freeze_quant_params is True since in GPTQ the activation quantization parameters should not be trained.
+        activation_quantizers.append(quantizer_class(cfg, freeze_quant_params=True))
 
     return weights_quantizers, activation_quantizers

model_compression_toolkit/gptq/keras/quantizer/soft_rounding/soft_quantizer_reg.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import List
+from typing import List, Callable
 
 import tensorflow as tf
 from keras import Model
@@ -22,61 +22,21 @@ from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribut
 from model_compression_toolkit.trainable_infrastructure import KerasTrainableQuantizationWrapper
 
 
-class LinearTempDecay:
-    """
-    Annealing process for the soft quantizer regularization temperature term.
-    """
-
-    def __init__(self, t_max: int, rel_start_decay: float = 0.2, start_b: int = 20, end_b: int = 2):
-        """
-        Initializes a LinearTempDecay object.
-
-        Args:
-            t_max: maximal time step.
-            rel_start_decay: Decay step size at the beginning of the process.
-            start_b: Starting value of the regularization term.
-            end_b: Target value of the regularization term.
-        """
-
-        self.t_max = t_max
-        self.start_decay = rel_start_decay * t_max
-        self.start_b = start_b
-        self.end_b = end_b
-
-    def __call__(self, t: int) -> float:
-        """
-        Cosine annealing scheduler for soft quantizer regularization temperature term.
-
-        Args:
-            t: The current time step.
-
-        Returns: Scheduled temperature.
-        """
-
-        is_before_start_decay = tf.cast(t < self.start_decay, tf.float32)
-
-        rel_t = (t - self.start_decay) / (self.t_max - self.start_decay)
-
-        return self.start_b * is_before_start_decay + \
-               (1 - is_before_start_decay) * \
-               (self.end_b + (self.start_b - self.end_b) * tf.math.maximum(0.0, (1 - rel_t)))
-
 
 class SoftQuantizerRegularization:
     """
     A class to handle the computation of soft quantizer regularization for GPTQ training.
     """
 
-    def __init__(self, total_gradient_steps: int):
+    def __init__(self, beta_scheduler: Callable[[int], float]):
         """
         Initializes the regularization computation object with a LinearDecay object.
 
         Args:
-            total_gradient_steps: The number of gradient steps during optimization.
+            beta_scheduler: a callable that accepts current time step and returns a corresponding beta value.
         """
         # Initializing the temperature decay according to the number of expected gradient steps
-        self.linear_decay = LinearTempDecay(total_gradient_steps)
-
+        self.beta_scheduler = beta_scheduler
         self.count_iter = tf.Variable(0.)
 
 
@@ -91,7 +51,7 @@ class SoftQuantizerRegularization:
         Returns: Regularization value.
         """
         soft_reg_aux: List[tf.Tensor] = []
-        b = self.linear_decay(self.count_iter.value())
+        b = self.beta_scheduler(self.count_iter.value())
         for layer in model.layers:
             if isinstance(layer, KerasTrainableQuantizationWrapper):
                 kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=type(layer.layer),

model_compression_toolkit/gptq/pytorch/gptq_training.py

@@ -17,15 +17,18 @@ from typing import Callable, List, Tuple, Union, Generator
 
 import numpy as np
 import torch
-from mct_quantizers import PytorchQuantizationWrapper, PytorchActivationQuantizationHolder
 from torch.nn import Module
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
+from model_compression_toolkit.gptq.common.gradual_activation_quantization import get_gradual_activation_quantizer_wrapper_factory
+from model_compression_toolkit.gptq.common.regularization_factory import get_regularization
+
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.hessian import HessianInfoService, HessianScoresGranularity
+
 from model_compression_toolkit.core.pytorch.back2framework.pytorch_model_builder import PyTorchModelBuilder
 from model_compression_toolkit.core.pytorch.constants import BIAS
 from model_compression_toolkit.core.pytorch.data_util import FixedDatasetFromGenerator, IterableDatasetFromGenerator, \
@@ -34,14 +37,15 @@ from model_compression_toolkit.core.pytorch.utils import to_torch_tensor, set_mo
 from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
 from model_compression_toolkit.gptq.common.gptq_training import GPTQTrainer
-from model_compression_toolkit.gptq.pytorch.graph_info import get_gptq_trainable_parameters, \
-    get_weights_for_loss
-from model_compression_toolkit.gptq.pytorch.quantizer.gradual_activation_quantization import \
-    get_gradual_activation_quantizer_wrapper_factory
+from model_compression_toolkit.gptq.pytorch.graph_info import get_gptq_trainable_parameters, get_weights_for_loss
 from model_compression_toolkit.gptq.pytorch.quantizer.quantization_builder import quantization_builder
-from model_compression_toolkit.gptq.pytorch.quantizer.regularization_factory import get_regularization
+
+from mct_quantizers import PytorchQuantizationWrapper, PytorchActivationQuantizationHolder
+from model_compression_toolkit.trainable_infrastructure.common.util import get_total_grad_steps
+from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import PytorchLinearAnnealingScheduler
+from model_compression_toolkit.gptq.pytorch.quantizer.soft_rounding.soft_quantizer_reg import SoftQuantizerRegularization as PytorchSoftQuantizerRegularization
+
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.trainable_infrastructure.pytorch.util import get_total_grad_steps
 
 
 class PytorchGPTQTrainer(GPTQTrainer):
@@ -78,7 +82,7 @@ class PytorchGPTQTrainer(GPTQTrainer):
 
         # must be set prior to model building in the base class constructor
         self.gradual_act_quantizer_wrapper_factory = get_gradual_activation_quantizer_wrapper_factory(
-            gptq_config, _get_total_grad_steps)
+            gptq_config, _get_total_grad_steps, PytorchLinearAnnealingScheduler)
 
         super().__init__(graph_float,
                          graph_quant,
@@ -121,7 +125,7 @@ class PytorchGPTQTrainer(GPTQTrainer):
         else:
             self.train_dataloader = self._prepare_train_dataloader_for_non_sla(representative_data_gen)
 
-        self.reg_func = get_regularization(self.gptq_config, _get_total_grad_steps)
+        self.reg_func = get_regularization(self.gptq_config, _get_total_grad_steps, PytorchSoftQuantizerRegularization, PytorchLinearAnnealingScheduler)
 
     def _prepare_train_dataloader_sla(self, data_gen_fn: Callable[[], Generator]) -> DataLoader:
         """

model_compression_toolkit/gptq/pytorch/quantization_facade.py

@@ -26,18 +26,15 @@ from model_compression_toolkit.core.common.visualization.tensorboard_writer impo
 from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.gptq.common.gptq_config import (
     GradientPTQConfig, GPTQHessianScoresConfig, GradualActivationQuantizationConfig)
-from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT
-from model_compression_toolkit.gptq.keras.quantization_facade import GPTQ_MOMENTUM
+from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT, LR_DEFAULT, LR_REST_DEFAULT, \
+    LR_BIAS_DEFAULT, GPTQ_MOMENTUM
 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.target_platform import TargetPlatformCapabilities
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 
-LR_DEFAULT = 1e-4
-LR_REST_DEFAULT = 1e-4
-LR_BIAS_DEFAULT = 1e-4
-LR_QUANTIZATION_PARAM_DEFAULT = 1e-4
+
 
 if FOUND_TORCH:
     from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
@@ -76,10 +73,6 @@ if FOUND_TORCH:
             regularization_factor (float): A floating point number that defines the regularization factor.
             hessian_batch_size (int): Batch size for Hessian computation in Hessian-based weights GPTQ.
             use_hessian_sample_attention (bool): whether to use Sample-Layer Attention score for weighted loss.
-            gradual_activation_quantization (bool, GradualActivationQuantizationConfig):
-              If False, GradualActivationQuantization is disabled.
-              If True, GradualActivationQuantization is enabled with the default settings.
-              GradualActivationQuantizationConfig object can be passed to use non-default settings.
 
         returns:
             a GradientPTQConfig object to use when fine-tuning the quantized model using gptq.
@@ -96,9 +89,6 @@ if FOUND_TORCH:
             >>> import torch
             >>> gptq_conf = mct.gptq.get_pytorch_gptq_config(n_epochs=3, optimizer=torch.optim.Adam([torch.Tensor(1)]))
 
-            To enable Gradual Activation Quantization with non-default settings build GradualActivationQuantizationConfig:
-            >>> gradual_act_conf = mct.gptq.GradualActivationQuantizationConfig(mct.gptq.QFractionLinearAnnealingConfig(initial_q_fraction=0.2))
-            >>> gptq_conf = mct.gptq.get_pytorch_gptq_config(n_epochs=3, gradual_activation_quantization=gradual_act_conf)
             The configuration can be passed to :func:`~model_compression_toolkit.pytorch_gradient_post_training_quantization` in order to quantize a pytorch model using gptq.
 
         """

model_compression_toolkit/qat/keras/quantizer/lsq/symmetric_lsq.py

@@ -18,7 +18,6 @@ from typing import Union
 import numpy as np
 import tensorflow as tf
 from tensorflow.python.framework.tensor_shape import TensorShape
-from model_compression_toolkit.constants import SIGNED
 
 from model_compression_toolkit.trainable_infrastructure import TrainingMethod
 
@@ -29,10 +28,9 @@ from model_compression_toolkit.qat.common import THRESHOLD_TENSOR
 from model_compression_toolkit import constants as C
 
 from model_compression_toolkit.qat.keras.quantizer.base_keras_qat_weight_quantizer import BaseKerasQATWeightTrainableQuantizer
-from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig
-from mct_quantizers.keras.quantizers import WeightsPOTInferableQuantizer, WeightsSymmetricInferableQuantizer, ActivationPOTInferableQuantizer, ActivationSymmetricInferableQuantizer
+from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig
+from mct_quantizers.keras.quantizers import WeightsPOTInferableQuantizer, WeightsSymmetricInferableQuantizer
 from model_compression_toolkit.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
-from model_compression_toolkit.qat.keras.quantizer.quant_utils import ste_round, grad_scale
 from model_compression_toolkit.trainable_infrastructure.keras.quantizer_utils import symmetric_lsq_quantizer
 
 

model_compression_toolkit/trainable_infrastructure/common/annealing_schedulers.py

@@ -0,0 +1,68 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from abc import abstractmethod
+
+
+class BaseLinearAnnealingScheduler:
+    def __init__(self, t_start: int, t_end: int, initial_val: float, target_val: float):
+        """
+        Base class for Linear annealing scheduler. Returns the corresponding annealed value per time step.
+
+        Args:
+            t_start: Time step to begin annealing.
+            t_end: Time step to complete annealing.
+            initial_val: Initial value before annealing.
+            target_val: Target value after annealing.
+
+        Raises:
+            ValueError: If t_start is not in the range [0, t_end).
+        """
+        if not (0 <= t_start < t_end):
+            raise ValueError(f'Expected 0 <= t_start < t_end, actual {t_end=} {t_start=}')
+
+        self.t_start = t_start
+        self.t_end = t_end
+        self.initial_val = initial_val
+        self.target_val = target_val
+
+    @abstractmethod
+    def _compute_factor(self, t: int) -> float:
+        """
+        Abstract method to compute the annealing factor based on time step `t`.
+
+        Args:
+            t: Current time step.
+
+        Returns:
+            float: Annealing factor, typically in the range [0, 1].
+
+        Raises:
+            NotImplementedError: If this method is not overridden in the subclass.
+        """
+        raise NotImplementedError("This method should be overridden in subclasses")
+
+    def __call__(self, t: int) -> float:
+        """
+        Calculates the annealed value based on the current time step `t`.
+
+        Args:
+            t: Current time step.
+
+        Returns:
+            float: Annealed value between initial_val and target_val.
+        """
+        factor = self._compute_factor(t)
+        return self.initial_val + factor * (self.target_val - self.initial_val)
+

model_compression_toolkit/trainable_infrastructure/keras/annealing_schedulers.py

@@ -0,0 +1,32 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import tensorflow as tf
+
+from model_compression_toolkit.trainable_infrastructure.common.annealing_schedulers import BaseLinearAnnealingScheduler
+
+
+class KerasLinearAnnealingScheduler(BaseLinearAnnealingScheduler):
+    def _compute_factor(self, t: int) -> float:
+        """
+        Computes the annealing factor for Keras models.
+
+        Args:
+            t: Current time step.
+
+        Returns:
+            float: Clipped annealing factor between 0 and 1.
+        """
+        factor = (t - self.t_start) / (self.t_end - self.t_start)
+        return tf.clip_by_value(factor, 0, 1)

model_compression_toolkit/trainable_infrastructure/pytorch/annealing_schedulers.py

@@ -13,27 +13,19 @@
 # limitations under the License.
 # ==============================================================================
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
+from model_compression_toolkit.trainable_infrastructure.common.annealing_schedulers import BaseLinearAnnealingScheduler
 
 
-class LinearAnnealingScheduler:
-    def __init__(self, t_start: int, t_end: int, initial_val: float, target_val: float):
+class PytorchLinearAnnealingScheduler(BaseLinearAnnealingScheduler):
+    def _compute_factor(self, t: int) -> float:
         """
-        Linear annealing scheduler. Returns the corresponding annealed value per time step.
+        Computes the annealing factor for torch models.
 
         Args:
-            t_start: time step to begin annealing.
-            t_end: time step to complete annealing.
-            initial_val: initial value.
-            target_val: target value.
-        """
-        if not (0 <= t_start < t_end):
-            raise ValueError(f'Expected 0 <= t_start < t_end, actual {t_end=} {t_start=}')
-
-        self.t_start = t_start
-        self.t_end = t_end
-        self.initial_val = initial_val
-        self.target_val = target_val
+            t: Current time step.
 
-    def __call__(self, t: int) -> float:
-        factor = to_torch_tensor((t - self.t_start) / (self.t_end - self.t_start)).clip(0, 1)
-        return self.initial_val + factor * (self.target_val - self.initial_val)
+        Returns:
+            float: Clipped annealing factor between 0 and 1.
+        """
+        factor = to_torch_tensor((t - self.t_start) / (self.t_end - self.t_start))
+        return factor.clip(0, 1)

tests_pytest/keras/gptq/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

tests_pytest/keras/gptq/test_gradual_act_quantization.py

@@ -0,0 +1,102 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from unittest.mock import Mock
+import pytest
+import numpy as np
+import tensorflow as tf
+
+from model_compression_toolkit.gptq.common.gradual_activation_quantization import GradualActivationQuantizerWrapper, \
+    get_gradual_activation_quantizer_wrapper_factory
+from model_compression_toolkit.trainable_infrastructure.keras.annealing_schedulers import KerasLinearAnnealingScheduler
+from model_compression_toolkit.gptq import GradientPTQConfig, GradualActivationQuantizationConfig, QFractionLinearAnnealingConfig
+
+
+
+@pytest.fixture
+def x():
+    return tf.random.normal((2, 5, 6, 7), seed=42, dtype=tf.float32)
+
+
+class Quantizer:
+    def __call__(self, x, training):
+        self.training = training
+        return 3 * x + 1
+
+
+class TestGradualActivationQuantization:
+
+    def test_gradual_act_quant_wrapper(self, x):
+        quantizer = Quantizer()
+        qw = GradualActivationQuantizerWrapper(quantizer, q_fraction_scheduler=lambda t: t / (t + 1))
+
+        y0, y1, y2 = [qw(x, training=True) for _ in range(3)]
+        assert np.allclose(y0.numpy(), x.numpy())  # t=0
+        assert np.allclose(y1.numpy(), 0.5 * x.numpy() + (1.5 * x.numpy() + 0.5))  # t=1
+        assert np.allclose(y2.numpy(), x.numpy() / 3 + (2 * x.numpy() + 2 / 3)) # t=2
+        assert quantizer.training is True
+
+        _ = qw(x, training=False)
+        assert quantizer.training is False  # correct flag was propagated
+
+    def test_factory_no_qdrop(self):
+        quantizer_wrapper, quantizer = self._run_factory_test(qdrop_cfg=None, get_grad_steps_fn=None)
+        assert quantizer_wrapper is quantizer
+
+    @pytest.mark.parametrize('end_step', (20, None))
+    def test_factory_linear(self, x, end_step):
+        qdrop_cfg = GradualActivationQuantizationConfig(
+            QFractionLinearAnnealingConfig(initial_q_fraction=0.3, target_q_fraction=0.8, start_step=10,
+                                           end_step=end_step)
+        )
+
+        def get_total_steps():
+            if end_step is None:
+                return 50
+            assert False  # should not be called if end_step is passed
+
+        quantizer_wrapper, quantizer = self._run_factory_test(qdrop_cfg, get_total_steps)
+
+        scheduler = quantizer_wrapper.q_fraction_scheduler
+        assert isinstance(scheduler, KerasLinearAnnealingScheduler)
+        exp_end_step = 50 if end_step is None else end_step
+        assert scheduler.t_start == 10
+        assert scheduler.t_end == exp_end_step
+        assert scheduler.initial_val == 0.3
+        assert scheduler.target_val == 0.8
+
+        y = [quantizer_wrapper(x, training=True) for _ in range(exp_end_step + 1)]
+
+        assert np.allclose(y[9].numpy(), 0.7 * x.numpy() + 0.3 * quantizer(x, training=True).numpy())
+        assert np.allclose(y[10].numpy(), 0.7 * x.numpy() + 0.3 * quantizer(x, training=True).numpy())
+        assert np.allclose(y[-1].numpy(), 0.2 * x.numpy() + 0.8 * quantizer(x, training=True).numpy())
+
+    def test_factory_linear_common_case(self, x):
+        # validate that we actually implemented the right thing - on first call float input, on last call fully quantized
+        qdrop_cfg = GradualActivationQuantizationConfig(
+            QFractionLinearAnnealingConfig(initial_q_fraction=0, target_q_fraction=1, start_step=0, end_step=None)
+        )
+        quantizer_wrapper, quantizer = self._run_factory_test(qdrop_cfg, lambda: 15)
+        y0, *_, y_last = [quantizer_wrapper(x, training=True) for _ in range(16)]
+        assert np.array_equal(y0.numpy(), x.numpy())
+        assert np.allclose(y_last.numpy(), quantizer(x, training=True).numpy())
+
+    def _run_factory_test(self, qdrop_cfg, get_grad_steps_fn):
+        # Mocks are used to just pass anything
+        gptq_cfg = GradientPTQConfig(n_epochs=5, optimizer=Mock(), loss=Mock(),
+                                     gradual_activation_quantization_config=qdrop_cfg)
+        factory = get_gradual_activation_quantizer_wrapper_factory(gptq_cfg, get_grad_steps_fn, KerasLinearAnnealingScheduler)
+        quantizer = Quantizer()
+        quantizer_wrapper = factory(quantizer)
+        return quantizer_wrapper, quantizer

tests_pytest/keras/trainable_infrastructure/__init__.py

@@ -0,0 +1,16 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+
+

tests_pytest/keras/trainable_infrastructure/test_linear_annealing.py

@@ -0,0 +1,49 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import numpy as np
+import pytest
+
+from model_compression_toolkit.trainable_infrastructure.keras.annealing_schedulers import KerasLinearAnnealingScheduler
+
+
+def test_linear_annealing():
+    scheduler = KerasLinearAnnealingScheduler(t_start=10, t_end=35, initial_val=3.4, target_val=-1.6)
+    for t in [0, 9, 10]:
+        assert _isclose(scheduler(t), 3.4)
+
+    for t in [35, 36, 1000]:
+        assert _isclose(scheduler(t), -1.6)
+
+    assert _isclose(scheduler(11), 3.2)
+    assert _isclose(scheduler(27), 0.)
+    assert _isclose(scheduler(34), -1.4)
+
+
+def test_linear_annealing_ascending():
+    scheduler = KerasLinearAnnealingScheduler(t_start=0, t_end=5, initial_val=-0.5, target_val=1.5)
+    assert _isclose(scheduler(0), -0.5)
+    assert _isclose(scheduler(1), -0.1)
+    assert _isclose(scheduler(4), 1.1)
+    assert _isclose(scheduler(5), 1.5)
+
+
+@pytest.mark.parametrize('start', [5, -1])
+def test_invalid(start):
+    with pytest.raises(ValueError):
+        KerasLinearAnnealingScheduler(t_start=start, t_end=4, initial_val=1, target_val=0)
+
+
+def _isclose(x, y):
+    return np.isclose(x, y)

tests_pytest/pytorch/gptq/test_gradual_act_quantization.py

@@ -18,9 +18,9 @@ import pytest
 import torch
 
 from model_compression_toolkit.core.pytorch.pytorch_device_config import get_working_device
-from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import LinearAnnealingScheduler
+from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import PytorchLinearAnnealingScheduler
 from model_compression_toolkit.gptq import GradientPTQConfig, GradualActivationQuantizationConfig, QFractionLinearAnnealingConfig
-from model_compression_toolkit.gptq.pytorch.quantizer.gradual_activation_quantization import (
+from model_compression_toolkit.gptq.common.gradual_activation_quantization import (
     GradualActivationQuantizerWrapper, get_gradual_activation_quantizer_wrapper_factory)
 
 
@@ -68,7 +68,7 @@ class TestGradualActivationQuantization:
         quantizer_wrapper, quantizer = self._run_factory_test(qdrop_cfg, get_total_steps)
 
         scheduler = quantizer_wrapper.q_fraction_scheduler
-        assert isinstance(scheduler, LinearAnnealingScheduler)
+        assert isinstance(scheduler, PytorchLinearAnnealingScheduler)
         exp_end_step = 50 if end_step is None else end_step
         assert scheduler.t_start == 10
         assert scheduler.t_end == exp_end_step
@@ -94,7 +94,7 @@ class TestGradualActivationQuantization:
         # Mocks are used to just pass anything
         gptq_cfg = GradientPTQConfig(n_epochs=5, optimizer=Mock(), loss=Mock(),
                                      gradual_activation_quantization_config=qdrop_cfg)
-        factory = get_gradual_activation_quantizer_wrapper_factory(gptq_cfg, get_grad_steps_fn)
+        factory = get_gradual_activation_quantizer_wrapper_factory(gptq_cfg, get_grad_steps_fn, PytorchLinearAnnealingScheduler)
         quantizer = Quantizer()
         quantizer_wrapper = factory(quantizer)
         return quantizer_wrapper, quantizer

tests_pytest/pytorch/trainable_infrastructure/test_linear_annealing.py

@@ -15,11 +15,11 @@
 import torch
 import pytest
 
-from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import LinearAnnealingScheduler
+from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import PytorchLinearAnnealingScheduler
 
 
 def test_linear_annealing():
-    scheduler = LinearAnnealingScheduler(t_start=10, t_end=35, initial_val=3.4, target_val=-1.6)
+    scheduler = PytorchLinearAnnealingScheduler(t_start=10, t_end=35, initial_val=3.4, target_val=-1.6)
     for t in [0, 9, 10]:
         assert _isclose(scheduler(t), 3.4)
 
@@ -32,7 +32,7 @@ def test_linear_annealing():
 
 
 def test_linear_annealing_ascending():
-    scheduler = LinearAnnealingScheduler(t_start=0, t_end=5, initial_val=-0.5, target_val=1.5)
+    scheduler = PytorchLinearAnnealingScheduler(t_start=0, t_end=5, initial_val=-0.5, target_val=1.5)
     assert _isclose(scheduler(0), -0.5)
     assert _isclose(scheduler(1), -0.1)
     assert _isclose(scheduler(4), 1.1)
@@ -42,7 +42,7 @@ def test_linear_annealing_ascending():
 @pytest.mark.parametrize('start', [5, -1])
 def test_invalid(start):
     with pytest.raises(ValueError):
-        LinearAnnealingScheduler(t_start=start, t_end=4, initial_val=1, target_val=0)
+        PytorchLinearAnnealingScheduler(t_start=start, t_end=4, initial_val=1, target_val=0)
 
 
 def _isclose(x, y):

model_compression_toolkit/gptq/keras/quantizer/regularization_factory.py

@@ -1,44 +0,0 @@
-# Copyright 2023 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-from tqdm import tqdm
-from typing import Callable
-
-from model_compression_toolkit.gptq import RoundingType, GradientPTQConfig, GradientPTQConfig
-from model_compression_toolkit.gptq.keras.quantizer.soft_rounding.soft_quantizer_reg import \
-    SoftQuantizerRegularization
-
-
-def get_regularization(gptq_config: GradientPTQConfig, representative_data_gen: Callable) -> Callable:
-    """
-    Returns a function that computes the regularization term for GPTQ training based on the given
-    rounding type in the GPTQ configuration.
-
-    Args:
-        gptq_config: A GPTQ configuration.
-        representative_data_gen: Dataset used for the GPTQ training.
-
-    Returns: A function for computing the regularization. If there is no regularization function defined for the given
-        rounding type, then it returns a function that just returns 0.
-
-    """
-    if gptq_config.rounding_type == RoundingType.SoftQuantizer:
-        # dry run on the representative dataset to count number of batches
-        num_batches = 0
-        for _ in tqdm(representative_data_gen(), "GPTQ initialization"):
-            num_batches += 1
-
-        return SoftQuantizerRegularization(total_gradient_steps=num_batches * gptq_config.n_epochs)
-    else:
-        return lambda m, e_reg: 0