mct-nightly 2.0.0.20240410.422__py3-none-any.whl → 2.0.0.20240412.408__py3-none-any.whl

{mct_nightly-2.0.0.20240410.422.dist-info → mct_nightly-2.0.0.20240412.408.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.0.0.20240410.422
+Version: 2.0.0.20240412.408
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN
@@ -23,7 +23,7 @@ Requires-Dist: PuLP
 Requires-Dist: matplotlib
 Requires-Dist: scipy
 Requires-Dist: protobuf
-Requires-Dist: mct-quantizers ==1.4
+Requires-Dist: mct-quantizers ==1.5
 
 # Model Compression Toolkit (MCT)
 

{mct_nightly-2.0.0.20240410.422.dist-info → mct_nightly-2.0.0.20240412.408.dist-info}/RECORD

@@ -1,7 +1,8 @@
-model_compression_toolkit/__init__.py,sha256=c33LV9Kt6hpVEoLixt_I5rqhtSzRBPSrdmFEifg-VHU,1573
-model_compression_toolkit/constants.py,sha256=KW_HUEPmQEYqCvWGyORqkYxpvO7w5LViB5J5D-pm_6o,3648
+model_compression_toolkit/__init__.py,sha256=EgLjEka0UxNeko2XoTELtIHjauH-j9CIyzQ86cFqdHg,1573
+model_compression_toolkit/constants.py,sha256=f9at1H_-vb5nvdHRmAHUco4ja4_QermK6yu0N9qbRGE,3723
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
+model_compression_toolkit/metadata.py,sha256=IyoON37lBv3TI0rZGCP4K5t3oYI4TOmYy-LRXOwHGpE,1136
 model_compression_toolkit/core/__init__.py,sha256=TrRgkWpT1AN2Faw1M_1HXyJkJnbxfn9p-RigDZl7pg0,1982
 model_compression_toolkit/core/analyzer.py,sha256=X-2ZpkH1xdXnISnw1yJvXnvV-ssoUh-9LkLISSWNqiY,3691
 model_compression_toolkit/core/graph_prep_runner.py,sha256=Ftqm59hT5TGWmSNkY9bFZkVfCacpGyZfCe-6yZR5WY0,10100
@@ -29,7 +30,7 @@ model_compression_toolkit/core/common/collectors/statistics_collector.py,sha256=
 model_compression_toolkit/core/common/fusion/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/core/common/fusion/layer_fusing.py,sha256=lOubqpc18TslhXZijWUJQAa1c3jIB2S-M-5HK78wJPQ,5548
 model_compression_toolkit/core/common/graph/__init__.py,sha256=Xr-Lt_qXMdrCnnOaUS_OJP_3iTTGfPCLf8_vSrQgCs0,773
-model_compression_toolkit/core/common/graph/base_graph.py,sha256=vS5mGRfVQWsFP28roJVSplVnpNkepacxGEuO7N-IPC0,38178
+model_compression_toolkit/core/common/graph/base_graph.py,sha256=06mvCb_HHA5iIOdQ31a-nimhrpSA-jYnuV1Ir76QGa8,38259
 model_compression_toolkit/core/common/graph/base_node.py,sha256=jPYpf6sci8LswatxTyygD8ZM5OvsCnxBEWsSl-g64wI,28492
 model_compression_toolkit/core/common/graph/edge.py,sha256=buoSEUZwilWBK3WeBKpJ-GeDaUA1SDdOHxDpxU_bGpk,3784
 model_compression_toolkit/core/common/graph/functional_node.py,sha256=RgwWAoMX7YV5c2gZdTBSX-ziTh3OLbebZXr3jitkxDs,3173
@@ -113,7 +114,7 @@ model_compression_toolkit/core/common/quantization/quantization_params_generatio
 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=BiwDqt5CeU6CW0Qusy3LwWhFtf2J9BvSuGMsTsG6rSw,8538
 model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_activations_computation.py,sha256=noEdvGiyyW7acgQ2OFWLedCODibTGYJifC9qo8YIU5U,4558
-model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_computation.py,sha256=wAeLTGsbMiUrkTrIdozWN8U5ZESSJzF1p0ZpPywVlw4,4346
+model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_computation.py,sha256=H2D9rdChIviL_j0mF6zy8Qeu_ZXKRu-hLqckSAT1MR8,4352
 model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_search.py,sha256=7kt0JB8PQE0SW9kg8fCwZ5mBkHNgiRrn0of4ZQYQN2A,41524
 model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_weights_computation.py,sha256=nug6XgsywxYf57XF_Tnt2xwdf0zLLsajiZKEblo4lFc,3882
 model_compression_toolkit/core/common/quantization/quantization_params_generation/symmetric_selection.py,sha256=QtSAtdAb7sTgtoe9L6DnMFO7rjkOtpzE9kD9xmG7eYM,9743
@@ -303,29 +304,29 @@ model_compression_toolkit/exporter/model_exporter/fw_agonstic/quantization_forma
 model_compression_toolkit/exporter/model_exporter/keras/__init__.py,sha256=uZ2RigbY9O2PJ0Il8wPpS_s7frgg9WUGd_SHeKGyl1A,699
 model_compression_toolkit/exporter/model_exporter/keras/base_keras_exporter.py,sha256=-wr2n0yRlmFixXBeZuxg6Rzlvz-ZFUX-PJgSXhgMrEo,1593
 model_compression_toolkit/exporter/model_exporter/keras/export_serialization_format.py,sha256=v_-rOsWDFI-3k8CoJIr-XzT7ny8WXpAMteWRWtTzaeg,963
-model_compression_toolkit/exporter/model_exporter/keras/fakely_quant_keras_exporter.py,sha256=Mss2KjVS4cpVH6fkQG5m4XJMr6QTKfmsD_DFit2qt_A,11481
-model_compression_toolkit/exporter/model_exporter/keras/fakely_quant_tflite_exporter.py,sha256=sqzqQ8US24WgDbg_FoP1NQBgqCbSVwrVTWrxcyY0nPA,3514
-model_compression_toolkit/exporter/model_exporter/keras/int8_tflite_exporter.py,sha256=nGtpDTeH5Tdp7sjyuXsy_9TPpijDYp4nkz366DUUJ0Q,8048
+model_compression_toolkit/exporter/model_exporter/keras/fakely_quant_keras_exporter.py,sha256=n_iXPwMomMVJTZH9M1WV7OJo11ppXOWkANu41fIlsjY,11702
+model_compression_toolkit/exporter/model_exporter/keras/fakely_quant_tflite_exporter.py,sha256=XoFGkVBikKh1BuxObrMLjfVLDIgy3X0rhmEl08CdJls,3727
+model_compression_toolkit/exporter/model_exporter/keras/int8_tflite_exporter.py,sha256=iTUXaia8XLJmmWdk4iiCah9sxeIyBJy42s9_EpuPhnw,8261
 model_compression_toolkit/exporter/model_exporter/keras/keras_export_facade.py,sha256=Q2hVl0dpH7hcVSVD9Y5BihtEGlHWrLk-_Y2RNPbfQTg,5750
 model_compression_toolkit/exporter/model_exporter/keras/mctq_keras_exporter.py,sha256=qXXkv3X_wb7t622EOHwXIxfGLGaDqh0T0y4UxREi4Bo,1976
 model_compression_toolkit/exporter/model_exporter/pytorch/__init__.py,sha256=uZ2RigbY9O2PJ0Il8wPpS_s7frgg9WUGd_SHeKGyl1A,699
 model_compression_toolkit/exporter/model_exporter/pytorch/base_pytorch_exporter.py,sha256=UPVkEUQCMZ4Lld6CRnEOPEmlfe5vcQZG0Q3FwRBodD4,4021
 model_compression_toolkit/exporter/model_exporter/pytorch/export_serialization_format.py,sha256=bPevy6OBqng41PqytBR55e6cBEuyrUS0H8dWX4zgjQ4,967
-model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py,sha256=TedEyQN9TQzRnb8afQVF4Kx28IeyEJRWiUloQxkBHAY,5136
+model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py,sha256=r2pOWFK-mSG8OzRiKGVOG4skzX0ZiM0eiRuBsL-ThoI,6067
 model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_torchscript_pytorch_exporter.py,sha256=ksWV2A-Njo-wAxQ_Ye2sLIZXBWJ_WNyjT7-qFFwvV2o,2897
-model_compression_toolkit/exporter/model_exporter/pytorch/pytorch_export_facade.py,sha256=iVKanLUALkK7WIDmNLpcvHu02EO47GM-T8iLr4srp-Q,6351
+model_compression_toolkit/exporter/model_exporter/pytorch/pytorch_export_facade.py,sha256=yz5dPMX5r1d9LJV4rYFS1pXqCbVUxvUmV4LELWcRinQ,6350
 model_compression_toolkit/exporter/model_wrapper/__init__.py,sha256=7CF2zvpTrIEm8qnbuHnLZyTZkwBBxV24V8QA0oxGbh0,1187
 model_compression_toolkit/exporter/model_wrapper/fw_agnostic/__init__.py,sha256=pKAdbTCFM_2BrZXUtTIw0ouKotrWwUDF_hP3rPwCM2k,696
 model_compression_toolkit/exporter/model_wrapper/fw_agnostic/get_inferable_quantizers.py,sha256=Bd3QhAR__YC9Xmobd5qHv9ofh_rPn_eTFV0sXizcBnY,2297
 model_compression_toolkit/exporter/model_wrapper/keras/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/exporter/model_wrapper/keras/validate_layer.py,sha256=YffgbVYJG5LKeIsW84Pi7NqzQcvJMeQRnAKQCCmIL6c,3776
 model_compression_toolkit/exporter/model_wrapper/keras/builder/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/exporter/model_wrapper/keras/builder/fully_quantized_model_builder.py,sha256=NBDzg2rX5BcVELtExHxS5wi0HFxwpGrEedB4ZPSVMas,5130
+model_compression_toolkit/exporter/model_wrapper/keras/builder/fully_quantized_model_builder.py,sha256=k3UrGAw6vKTmZ-oO1lv0VqK3IpAiet9jlIHyEIoL2u0,5132
 model_compression_toolkit/exporter/model_wrapper/keras/builder/node_to_quantizer.py,sha256=uL6tJWC4s2IWUy8GJVwtMWpwZZioRRztfKyPJHo14xI,9442
 model_compression_toolkit/exporter/model_wrapper/pytorch/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/exporter/model_wrapper/pytorch/validate_layer.py,sha256=uTQcnzvP44CgPO0twsUdiMmTBE_Td6ZdQtz5U0GZuPI,3464
 model_compression_toolkit/exporter/model_wrapper/pytorch/builder/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/exporter/model_wrapper/pytorch/builder/fully_quantized_model_builder.py,sha256=T3QNZl0JFRAm62Z66quHPx0iNHgXwyfSpoBgbqJBBnY,4915
+model_compression_toolkit/exporter/model_wrapper/pytorch/builder/fully_quantized_model_builder.py,sha256=D_mEUK1sb4kY5946oErfw3RC5mfBTVaw3LZRIKWYKcE,4918
 model_compression_toolkit/exporter/model_wrapper/pytorch/builder/node_to_quantizer.py,sha256=4sN5z-6BXrTE5Dp2FX_jKO9ty5iZ2r4RM7XvXtDVLSI,9348
 model_compression_toolkit/gptq/__init__.py,sha256=YKg-tMj9D4Yd0xW9VRD5EN1J5JrmlRbNEF2fOSgodqA,1228
 model_compression_toolkit/gptq/runner.py,sha256=MIg-oBtR1nbHkexySdCJD_XfjRoHSknLotmGBMuD5qM,5924
@@ -338,14 +339,14 @@ model_compression_toolkit/gptq/common/gptq_training.py,sha256=rLA1xlOO-6gWfmc2dL
 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=OhYfH6zxRHrRhCde0lbcV9Hu2oeDD9RXh-O8vOPgLbs,18875
+model_compression_toolkit/gptq/keras/gptq_training.py,sha256=zyVcEQzdnNsrIz32U1pqqoi08hzxRdJ2CumaPFGwbDM,19123
 model_compression_toolkit/gptq/keras/graph_info.py,sha256=5IvgGlJlgOmQYmldjdCBv7tuzAoY0HazatG5Pedrg0Q,4639
-model_compression_toolkit/gptq/keras/quantization_facade.py,sha256=zAkzWpWP9_aobWgMo_BlUm7-4fR5dHvoGx0sDqs2rZg,14299
+model_compression_toolkit/gptq/keras/quantization_facade.py,sha256=CCV9uyaq-qUGDeXL5OgEWFXSiUkerXrNwFVyA1brrKM,14663
 model_compression_toolkit/gptq/keras/quantizer/__init__.py,sha256=-DK1CDXvlsnEbki4lukZLpl6Xrbo91_jcqxXlG5Eg6Q,963
 model_compression_toolkit/gptq/keras/quantizer/base_keras_gptq_quantizer.py,sha256=2YU-x4-Q5f6hkUJf0tw6vcwdNwRMHdefrFjhhyHYsvA,4782
 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=7kvQQz2zHTRkIzJpsOPe8PWtfsOpcGZ2hjVIxbc-qJo,1906
+model_compression_toolkit/gptq/keras/quantizer/regularization_factory.py,sha256=guf7ygnLsZeWnTDz4yJdE2iTkd1oE0uQAZwKnGV3OAk,1957
 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/symmetric_soft_quantizer.py,sha256=BBSDWLmeywjSM5N6oJkMgcuo7zrXTesB4zLwRGG8QB0,12159
@@ -355,14 +356,14 @@ 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=kDuWw-6zh17wZpYWh4Xa94rpoodf82DksgjQCnL7nBc,2719
 model_compression_toolkit/gptq/pytorch/gptq_pytorch_implementation.py,sha256=tECPTavxn8EEwgLaP2zvxdJH6Vg9jC0YOIMJ7857Sdc,1268
-model_compression_toolkit/gptq/pytorch/gptq_training.py,sha256=LN4vOwcMuSSFTSnHDACV9hX_Yd2YIXJRl7WkdODuA0k,16245
+model_compression_toolkit/gptq/pytorch/gptq_training.py,sha256=xkDa62AdIRwv8dEshffALW9Ri66eseEpyUF9taMUKns,16509
 model_compression_toolkit/gptq/pytorch/graph_info.py,sha256=yXJzDd24zfGs2_vfMovxD1WSh1RxXoPxN4GztOf3P5c,3967
-model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=-4USg-tep6EQSArcTxBowhMeAuExrBTNLOWgHFpsIy4,12699
+model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=iBLEbLgde6JQNPhJysfT2rl_Sc7-wyoIZnXRAXQWnR0,13065
 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=TCA1hAc7raPnrjl06sjFtVM4XUtLtuwAhCGX4U3KGZo,4137
 model_compression_toolkit/gptq/pytorch/quantizer/quant_utils.py,sha256=OocYYRqvl7rZ37QT0hTzfJnWGiNCPskg7cziTlR7TRk,3893
 model_compression_toolkit/gptq/pytorch/quantizer/quantization_builder.py,sha256=uT9N_aBj965hvQfKd67fS1B0SXGnOLVcqa3wW4b2iZE,4566
-model_compression_toolkit/gptq/pytorch/quantizer/regularization_factory.py,sha256=-6fn6U6y2HZXluOfShYLeFKiuiDMVvsF64OTUDCrne4,1908
+model_compression_toolkit/gptq/pytorch/quantizer/regularization_factory.py,sha256=mDWZERLwtDzqWeJUwHMVyGdlS8wPLjJ3NvZiKBP6BNA,1959
 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=oO7WgsAHMnWoXNm_gTKAAe-Nd79mGL_m677ai-ui424,4132
 model_compression_toolkit/gptq/pytorch/quantizer/soft_rounding/symmetric_soft_quantizer.py,sha256=kLVQC1hXzDpP4Jx7AwnA764oGnY5AMEuvUUhAvhz09M,12347
@@ -377,9 +378,9 @@ model_compression_toolkit/pruning/pytorch/pruning_facade.py,sha256=cSuvHHCqgr7k9
 model_compression_toolkit/ptq/__init__.py,sha256=Z_hkmTh7aLFei1DJKV0oNVUbrv_Q_0CTw-qD85Xf8UM,904
 model_compression_toolkit/ptq/runner.py,sha256=_c1dSjlPPpsx59Vbg1buhG9bZq__OORz1VlPkwjJzoc,2552
 model_compression_toolkit/ptq/keras/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/ptq/keras/quantization_facade.py,sha256=T1_UqXmOc4I2a6IHkQAlFhGtcAYjsXSApMIdRlvgDvg,10154
+model_compression_toolkit/ptq/keras/quantization_facade.py,sha256=s6vBCK98l-R12yWASkutPSmNSfPX7457DazroJwhjpo,10517
 model_compression_toolkit/ptq/pytorch/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/ptq/pytorch/quantization_facade.py,sha256=eof9bo-Mv_lLY7fFpiVeT5pIde-MuTWkIAqRKH4j9MI,8646
+model_compression_toolkit/ptq/pytorch/quantization_facade.py,sha256=kkdgBXRBkblBTOW5EaySI_bN4_becSUwbdgOTb7FW2c,9012
 model_compression_toolkit/qat/__init__.py,sha256=kj2qsZh_Ca7PncsHKcaL5EVT2H8g4hYtvaQ3KFxOkwE,1143
 model_compression_toolkit/qat/common/__init__.py,sha256=6tLZ4R4pYP6QVztLVQC_jik2nES3l4uhML0qUxZrezk,829
 model_compression_toolkit/qat/common/qat_config.py,sha256=zoq0Vb74vCY7WlWD8JH_KPrHDoUHSvMc3gcO53u7L2U,3394
@@ -415,19 +416,19 @@ model_compression_toolkit/target_platform_capabilities/target_platform/current_t
 model_compression_toolkit/target_platform_capabilities/target_platform/fusing.py,sha256=f3xBAI6ivPvEj4lw8cAvTKdIbs7CRdLAa_0LvhGw3Dg,3924
 model_compression_toolkit/target_platform_capabilities/target_platform/op_quantization_config.py,sha256=xMBarITs2g_pCakep1vlNpM9Au5kQubTtFQAFJhvtYs,14248
 model_compression_toolkit/target_platform_capabilities/target_platform/operators.py,sha256=rRmrmPBY4rxCWVpEc6FxeOPUFh8MkfwgQsqD82U9a7w,3108
-model_compression_toolkit/target_platform_capabilities/target_platform/target_platform_model.py,sha256=u6xlc5sSsCYY02OBHXSOynlv7SW0bbyQB53L6a61LnM,9320
+model_compression_toolkit/target_platform_capabilities/target_platform/target_platform_model.py,sha256=55EQyP4McJBwsw9_l9eseEZ-V2ygQfdRXjjJLF6HdIM,9482
 model_compression_toolkit/target_platform_capabilities/target_platform/target_platform_model_component.py,sha256=TDbNQwmF7Id-FoIQZlR7ZOcz_nRb4XKBmDihAgKT0u8,1392
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/__init__.py,sha256=WCP1wfFZgM4eFm-pPeUinr5R_aSx5qwfSQqLZCXUNBA,1513
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/attribute_filter.py,sha256=jfhszvuD2Fyy6W2KjlLzXBQKFzTqGAaDZeFVr4-ONQw,8776
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/current_tpc.py,sha256=fIheShGOnxWYKqT8saHpBJqOU5RG_1Hp9qHry7IviIw,2115
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/layer_filter_params.py,sha256=Cl6-mACpje2jM8RJkibbqE3hvTkFR3r26-lW021mIiA,4019
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/operations_to_layers.py,sha256=1JN3yvNiJyDfva0tLTH3ej_qORzrQcPz32bSMKl49_0,6720
-model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/target_platform_capabilities.py,sha256=HK0KsFh0l94vl34N0tkHO2bckvCVO_xZMIl7A0bXGuM,9837
+model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/target_platform_capabilities.py,sha256=KP8IWlHzkXzVjqIiRtAW6sTYyHJ2wVFFX4hMt_N6o3s,9910
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/target_platform_capabilities_component.py,sha256=FvrYI0Qy7DCmDp2gyUYyCZq5pY84JgLtJqSIiVTJ8Ss,1030
 model_compression_toolkit/target_platform_capabilities/tpc_models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 model_compression_toolkit/target_platform_capabilities/tpc_models/get_target_platform_capabilities.py,sha256=aHoAu5Iye9YVn2HLwNb4X9cUDX1WJt20R5GsNGIAk9E,3337
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/__init__.py,sha256=lNJ29DYxaLUPDstRDA1PGI5r9Fulq_hvrZMlhst1Z5g,697
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/target_platform_capabilities.py,sha256=1O6q93V268cqUBq2tACRP8ZMk2hm7QjmDmrPubDW-jU,2959
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/target_platform_capabilities.py,sha256=fPOzybGECCWPkAD1hmJryWZrf9vd5Od-UOH6PE0lH94,3820
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/latest/__init__.py,sha256=F5RG4MnuAwKcNXbfVbPFLQu30-lNax-7knqu20B6udQ,1522
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1/__init__.py,sha256=1mMOREEMoNHu_KTMGDp4crN61opKWX6aFn1DrDLvqcc,717
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1/tp_model.py,sha256=S-GwMI-JiuPpbtOdd6TSOEjiUFiIs6M2RAiJNJ3O950,10883
@@ -441,6 +442,14 @@ model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_pot/tp_model.py,sha256=o1KloA8WPw1MbtZ-4p-kxQuroBAL67z77dPpliZyH9o,10369
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_pot/tpc_keras.py,sha256=NkAGCZbSgXYeRAiJRzt19h2cxkrVQJaHu8-2jHZLOYg,6505
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_pot/tpc_pytorch.py,sha256=X853xDEF-3rcPoqxbrlYN28vvW3buSdM36c_eN_LKx8,5758
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2/__init__.py,sha256=vKWAoQ2KkhuptS5HZB50zHG6KY8wHpHTxPugw_nGCRo,717
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2/tp_model.py,sha256=8FZjOCaQRwrQLbtmzNrrRj2-VyZMUGzsIWKIDpGVEoQ,10947
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2/tpc_keras.py,sha256=waefIjxpRfjSnJhKfgpR2DZM6B9NLzU-harUerk6oBc,6485
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2/tpc_pytorch.py,sha256=m6gOYBJZZd31RZHaiovDPLMlpt0HRiJhcKmk73_8380,5732
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_lut/__init__.py,sha256=wUk4Xsg7jpxOWYjq2K3WUwLcI185p_sVPK-ttG0ydhA,721
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_lut/tp_model.py,sha256=T6Hp_Rk15SAz22g_SWDyHJecBpBAjxKt3ezuVEYf4LE,10680
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_lut/tpc_keras.py,sha256=O7h77xAC92nlpgl6cwuJvL3DKDVBplMbKZEQXXg6p0w,6493
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_lut/tpc_pytorch.py,sha256=uvHIwmZ1DtA9-XY7SV-b4pSWZde2Ya9MNvGRlVVvdY0,5739
 model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/target_platform_capabilities.py,sha256=lnhJcwvTF0t7ybeiTleIS1p0aD8xzFZxVPx4ISk5uWQ,2090
 model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/latest/__init__.py,sha256=UUvUCcTots_sehdRnDfgkaE8WPQ7dPbeuhDF4Qy2nzw,1510
@@ -471,8 +480,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=7bbzqJN8ZAycVDvZr_5xC-niTAR5df8f03Kooev_pfg,3047
-mct_nightly-2.0.0.20240410.422.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.0.0.20240410.422.dist-info/METADATA,sha256=Xx2HTbZkpp4O8bS07IXSnaYSh9ZZTxe61I47ovv9fzE,18795
-mct_nightly-2.0.0.20240410.422.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-mct_nightly-2.0.0.20240410.422.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.0.0.20240410.422.dist-info/RECORD,,
+mct_nightly-2.0.0.20240412.408.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.0.0.20240412.408.dist-info/METADATA,sha256=7hBFTvvDA-m2U-nbakxLvO7VES7x2VTln6ZpE6JGn68,18795
+mct_nightly-2.0.0.20240412.408.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+mct_nightly-2.0.0.20240412.408.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.0.0.20240412.408.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.0.0.20240410.000422"
+__version__ = "2.0.0.20240412.000408"

model_compression_toolkit/constants.py

@@ -24,6 +24,10 @@ FOUND_ONNX = importlib.util.find_spec("onnx") is not None
 FOUND_ONNXRUNTIME = importlib.util.find_spec("onnxruntime") is not None
 FOUND_SONY_CUSTOM_LAYERS = importlib.util.find_spec('sony_custom_layers') is not None
 
+# Metadata fields
+MCT_VERSION = 'mct_version'
+TPC_VERSION = 'tpc_version'
+
 WEIGHTS_SIGNED = True
 # Minimal threshold to use for quantization ranges:
 MIN_THRESHOLD = (2 ** -16)

model_compression_toolkit/core/common/graph/base_graph.py

@@ -103,9 +103,10 @@ class Graph(nx.MultiDiGraph, GraphSearches):
             if n.is_custom:
                 if not is_node_in_tpc:
                     Logger.critical(f'MCT does not support optimizing Keras custom layers. Found a layer of type {n.type}. '
-                                 f' Please add the custom layer to Target Platform Capabilities (TPC), or file a feature request or an issue if you believe this should be supported.')
+                                    ' Please add the custom layer to Target Platform Capabilities (TPC), or file a feature '
+                                    'request or an issue if you believe this should be supported.')  # pragma: no cover
                 if any([qc.default_weight_attr_config.enable_weights_quantization for qc in n.get_qco(tpc).quantization_config_list]):
-                    Logger.critical(f'Layer identified: {n.type}. MCT does not support weight quantization for Keras custom layers.')
+                    Logger.critical(f'Layer identified: {n.type}. MCT does not support weight quantization for Keras custom layers.')  # pragma: no cover
 
         self.tpc = tpc
 

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

@@ -42,14 +42,14 @@ def calculate_quantization_params(graph: Graph,
 
     """
 
-    Logger.info(f"Running quantization parameters search. "
+    Logger.info(f"\nRunning quantization parameters search. "
                 f"This process might take some time, "
                 f"depending on the model size and the selected quantization methods.\n")
 
     # Create a list of nodes to compute their thresholds
     nodes_list: List[BaseNode] = nodes if specific_nodes else graph.nodes()
 
-    for n in tqdm(nodes_list, "Calculating quantization params"):  # iterate only nodes that we should compute their thresholds
+    for n in tqdm(nodes_list, "Calculating quantization parameters"):  # iterate only nodes that we should compute their thresholds
         for candidate_qc in n.candidates_quantization_cfg:
             for attr in n.get_node_weights_attributes():
                 if n.is_weights_quantization_enabled(attr):

model_compression_toolkit/exporter/model_exporter/keras/fakely_quant_keras_exporter.py

@@ -100,7 +100,7 @@ class FakelyQuantKerasExporter(BaseKerasExporter):
                         weights_list.append(layer.get_quantized_weights()['kernel'])
                     else:
                         Logger.critical(f'KerasQuantizationWrapper should wrap only DepthwiseConv2D, Conv2D, Dense'
-                                     f' and Conv2DTranspose layers but wrapped layer is {layer.layer}')
+                                        f' and Conv2DTranspose layers but wrapped layer is {layer.layer}')
 
                     if layer.layer.bias is not None:
                         weights_list.append(layer.layer.bias)
@@ -121,6 +121,11 @@ class FakelyQuantKerasExporter(BaseKerasExporter):
 
             return layer
 
+        # Delete metadata layer if exists
+        if hasattr(self.model, 'metadata_layer'):
+            Logger.info('Metadata is not exported to FakelyQuant models.')
+            delattr(self.model, 'metadata_layer')
+
         # clone each layer in the model and apply _unwrap_quantize_wrapper to layers wrapped with a QuantizeWrapper.
         self.exported_model = tf.keras.models.clone_model(self.model,
                                                           input_tensors=None,

model_compression_toolkit/exporter/model_exporter/keras/fakely_quant_tflite_exporter.py

@@ -56,6 +56,11 @@ class FakelyQuantTFLiteExporter(FakelyQuantKerasExporter):
         (namely, weights that are in fake-quant format) and fake-quant layers for the activations.
 
         """
+        # Delete metadata layer if exists
+        if hasattr(self.model, 'metadata_layer'):
+            Logger.info('Metadata is not exported to TFLite models.')
+            delattr(self.model, 'metadata_layer')
+
         # Use Keras exporter to quantize model's weights before converting it to TFLite.
         # Since exporter saves the model, we use a tmp path for saving, and then we delete it.
         handle, tmp_file = tempfile.mkstemp(DEFAULT_KERAS_EXPORT_EXTENTION)

model_compression_toolkit/exporter/model_exporter/keras/int8_tflite_exporter.py

@@ -166,6 +166,11 @@ class INT8TFLiteExporter(FakelyQuantKerasExporter):
 
             return layer_to_substitue
 
+        # Delete metadata layer if exists
+        if hasattr(self.model, 'metadata_layer'):
+            Logger.info('Metadata is not exported to TFLite models.')
+            delattr(self.model, 'metadata_layer')
+
         # Transform the model to a new model that can be converted to int8 models.
         # For example: replace dense layers with point-wise layers (to support per-channel quantization)
         self.transformed_model = clone_model(self.model,

model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py

@@ -13,17 +13,21 @@
 # limitations under the License.
 # ==============================================================================
 from typing import Callable
+from io import BytesIO
 
 import torch.nn
+import onnx
 
 from mct_quantizers import PytorchActivationQuantizationHolder, PytorchQuantizationWrapper
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
 from model_compression_toolkit.exporter.model_exporter.pytorch.base_pytorch_exporter import BasePyTorchExporter
 from mct_quantizers import pytorch_quantizers
+from mct_quantizers.pytorch.metadata import add_onnx_metadata
 
 DEFAULT_ONNX_OPSET_VERSION=15
 
+
 class FakelyQuantONNXPyTorchExporter(BasePyTorchExporter):
     """
     Exporter for fakely-quant PyTorch models.
@@ -58,7 +62,6 @@ class FakelyQuantONNXPyTorchExporter(BasePyTorchExporter):
         self._use_onnx_custom_quantizer_ops = use_onnx_custom_quantizer_ops
         self._onnx_opset_version = onnx_opset_version
 
-
     def export(self) -> None:
         """
         Convert an exportable (fully-quantized) PyTorch model to a fakely-quant model
@@ -74,7 +77,7 @@ class FakelyQuantONNXPyTorchExporter(BasePyTorchExporter):
         # If _use_onnx_custom_quantizer_ops is set to True, the quantizer forward function will use
         # the custom implementation when exporting the operator into onnx model. If not, it removes the
         # wraps and quantizes the ops in place (for weights, for activation torch quantization function is
-        # exported since it's used during forward.
+        # exported since it's used during forward).
         if self._use_onnx_custom_quantizer_ops:
             self._enable_onnx_custom_ops_export()
         else:
@@ -87,15 +90,30 @@ class FakelyQuantONNXPyTorchExporter(BasePyTorchExporter):
 
         model_input = to_torch_tensor(next(self.repr_dataset())[0])
 
-        torch.onnx.export(self.model,
-                          model_input,
-                          self.save_model_path,
-                          opset_version=self._onnx_opset_version,
-                          verbose=False,
-                          input_names=['input'],
-                          output_names=['output'],
-                          dynamic_axes={'input': {0: 'batch_size'},
-                                        'output': {0: 'batch_size'}})
+        if hasattr(self.model, 'metadata'):
+            onnx_bytes = BytesIO()
+            torch.onnx.export(self.model,
+                              model_input,
+                              onnx_bytes,
+                              opset_version=self._onnx_opset_version,
+                              verbose=False,
+                              input_names=['input'],
+                              output_names=['output'],
+                              dynamic_axes={'input': {0: 'batch_size'},
+                                            'output': {0: 'batch_size'}})
+            onnx_model = onnx.load_from_string(onnx_bytes.getvalue())
+            onnx_model = add_onnx_metadata(onnx_model, self.model.metadata)
+            onnx.save_model(onnx_model, self.save_model_path)
+        else:
+            torch.onnx.export(self.model,
+                              model_input,
+                              self.save_model_path,
+                              opset_version=self._onnx_opset_version,
+                              verbose=False,
+                              input_names=['input'],
+                              output_names=['output'],
+                              dynamic_axes={'input': {0: 'batch_size'},
+                                            'output': {0: 'batch_size'}})
 
     def _enable_onnx_custom_ops_export(self):
         """

model_compression_toolkit/exporter/model_exporter/pytorch/pytorch_export_facade.py

@@ -40,7 +40,7 @@ if FOUND_TORCH:
                              repr_dataset: Callable,
                              is_layer_exportable_fn: Callable = is_pytorch_layer_exportable,
                              serialization_format: PytorchExportSerializationFormat = PytorchExportSerializationFormat.ONNX,
-                             quantization_format : QuantizationFormat = QuantizationFormat.MCTQ,
+                             quantization_format: QuantizationFormat = QuantizationFormat.MCTQ,
                              onnx_opset_version=DEFAULT_ONNX_OPSET_VERSION) -> None:
         """
         Export a PyTorch quantized model to a torchscript or onnx model.

model_compression_toolkit/exporter/model_wrapper/keras/builder/fully_quantized_model_builder.py

@@ -90,7 +90,7 @@ if FOUND_TF:
                                                                                         fw_impl=C.keras.keras_implementation.KerasImplementation())).build_model()
         exportable_model.trainable = False
 
-        Logger.info("Please run your accuracy evaluation on the exported quantized model to verify it's accuracy.\n"
+        Logger.info("\nPlease run your accuracy evaluation on the exported quantized model to verify it's accuracy.\n"
                     "Checkout the FAQ and Troubleshooting pages for resolving common issues and improving the quantized model accuracy:\n"
                     "FAQ: https://github.com/sony/model_optimization/tree/main/FAQ.md\n"
                     "Quantization Troubleshooting: https://github.com/sony/model_optimization/tree/main/quantization_troubleshooting.md")

model_compression_toolkit/exporter/model_wrapper/pytorch/builder/fully_quantized_model_builder.py

@@ -82,7 +82,7 @@ if FOUND_TORCH:
                                                           get_activation_quantizer_holder(n,
                                                                                           fw_impl=C.pytorch.pytorch_implementation.PytorchImplementation())).build_model()
 
-        Logger.info("Please run your accuracy evaluation on the exported quantized model to verify it's accuracy.\n"
+        Logger.info("\nPlease run your accuracy evaluation on the exported quantized model to verify it's accuracy.\n"
                     "Checkout the FAQ and Troubleshooting pages for resolving common issues and improving the quantized model accuracy:\n"
                     "FAQ: https://github.com/sony/model_optimization/tree/main/FAQ.md\n"
                     "Quantization Troubleshooting: https://github.com/sony/model_optimization/tree/main/quantization_troubleshooting.md")
@@ -93,4 +93,4 @@ if FOUND_TORCH:
 else:
     def get_exportable_pytorch_model(*args, **kwargs):
         Logger.critical("PyTorch must be installed to use 'get_exportable_pytorch_model'. "
-                        "The 'torch' package is missing.")  # pragma: no cover
+                        "The 'torch' package is missing.")  # pragma: no cover

model_compression_toolkit/gptq/keras/gptq_training.py

@@ -301,21 +301,23 @@ class KerasGPTQTrainer(GPTQTrainer):
         Returns: None
 
         """
-        for _ in tqdm(range(n_epochs)):
-            for data in tqdm(data_function()):
-                input_data = [d * self.input_scale for d in data]
-
-                loss_value_step, grads = self.nano_training_step(input_data, in_compute_gradients,
-                                                                 in_optimizer_with_param, is_training)
-                # Run one step of gradient descent by updating
-                # the value of the variables to minimize the loss.
-                for i, (o, p) in enumerate(in_optimizer_with_param):
-                    o.apply_gradients(zip(grads[i], p))
-                if self.gptq_config.log_function is not None:
-                    self.gptq_config.log_function(loss_value_step, grads[0], in_optimizer_with_param[0][-1],
-                                                  self.compare_points)
-                self.loss_list.append(loss_value_step.numpy())
-                Logger.debug(f'last loss value: {self.loss_list[-1]}')
+        with tqdm(range(n_epochs), "Running GPTQ optimization") as epochs_pbar:
+            for _ in epochs_pbar:
+                with tqdm(data_function(), position=1, leave=False) as data_pbar:
+                    for data in data_pbar:
+                        input_data = [d * self.input_scale for d in data]
+
+                        loss_value_step, grads = self.nano_training_step(input_data, in_compute_gradients,
+                                                                         in_optimizer_with_param, is_training)
+                        # Run one step of gradient descent by updating
+                        # the value of the variables to minimize the loss.
+                        for i, (o, p) in enumerate(in_optimizer_with_param):
+                            o.apply_gradients(zip(grads[i], p))
+                        if self.gptq_config.log_function is not None:
+                            self.gptq_config.log_function(loss_value_step, grads[0], in_optimizer_with_param[0][-1],
+                                                          self.compare_points)
+                        self.loss_list.append(loss_value_step.numpy())
+                        Logger.debug(f'last loss value: {self.loss_list[-1]}')
 
     def update_graph(self):
         """

model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -31,6 +31,7 @@ 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
 
 LR_DEFAULT = 0.15
 LR_REST_DEFAULT = 1e-4
@@ -48,6 +49,7 @@ if FOUND_TF:
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from model_compression_toolkit.exporter.model_wrapper import get_exportable_keras_model
     from model_compression_toolkit import get_target_platform_capabilities
+    from mct_quantizers.keras.metadata import add_metadata
 
     # As from TF2.9 optimizers package is changed
     if version.parse(tf.__version__) < version.parse("2.9"):
@@ -234,7 +236,10 @@ if FOUND_TF:
                                         fw_impl,
                                         DEFAULT_KERAS_INFO)
 
-        return get_exportable_keras_model(tg_gptq)
+        exportable_model, user_info = get_exportable_keras_model(tg_gptq)
+        if target_platform_capabilities.tp_model.add_metadata:
+            exportable_model = add_metadata(exportable_model, get_versions_dict(target_platform_capabilities))
+        return exportable_model, user_info
 
 else:
     # If tensorflow is not installed,

model_compression_toolkit/gptq/keras/quantizer/regularization_factory.py

@@ -12,6 +12,7 @@
 # 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
@@ -35,7 +36,7 @@ def get_regularization(gptq_config: GradientPTQConfig, representative_data_gen:
     if gptq_config.rounding_type == RoundingType.SoftQuantizer:
         # dry run on the representative dataset to count number of batches
         num_batches = 0
-        for _ in representative_data_gen():
+        for _ in tqdm(representative_data_gen(), "GPTQ initialization"):
             num_batches += 1
 
         return SoftQuantizerRegularization(total_gradient_steps=num_batches * gptq_config.n_epochs)

model_compression_toolkit/gptq/pytorch/gptq_training.py

@@ -248,22 +248,24 @@ class PytorchGPTQTrainer(GPTQTrainer):
             data_function: A callable function that give a batch of samples.
             n_epochs: Number of update iterations of representative dataset.
         """
-        for _ in tqdm(range(n_epochs)):
-            for data in tqdm(data_function()):
-                input_data = [d * self.input_scale for d in data]
-                input_tensor = to_torch_tensor(input_data)
-                y_float = self.float_model(input_tensor)  # running float model
-                loss_value, grads = self.compute_gradients(y_float, input_tensor)
-                # Run one step of gradient descent by updating the value of the variables to minimize the loss.
-                for (optimizer, _) in self.optimizer_with_param:
-                    optimizer.step()
-                    optimizer.zero_grad()
-                if self.gptq_config.log_function is not None:
-                    self.gptq_config.log_function(loss_value.item(),
-                                                  torch_tensor_to_numpy(grads),
-                                                  torch_tensor_to_numpy(self.optimizer_with_param[0][-1]))
-                self.loss_list.append(loss_value.item())
-                Logger.debug(f'last loss value: {self.loss_list[-1]}')
+        with tqdm(range(n_epochs), "Running GPTQ optimization") as epochs_pbar:
+            for _ in epochs_pbar:
+                with tqdm(data_function(), position=1, leave=False) as data_pbar:
+                    for data in data_pbar:
+                        input_data = [d * self.input_scale for d in data]
+                        input_tensor = to_torch_tensor(input_data)
+                        y_float = self.float_model(input_tensor)  # running float model
+                        loss_value, grads = self.compute_gradients(y_float, input_tensor)
+                        # Run one step of gradient descent by updating the value of the variables to minimize the loss.
+                        for (optimizer, _) in self.optimizer_with_param:
+                            optimizer.step()
+                            optimizer.zero_grad()
+                        if self.gptq_config.log_function is not None:
+                            self.gptq_config.log_function(loss_value.item(),
+                                                          torch_tensor_to_numpy(grads),
+                                                          torch_tensor_to_numpy(self.optimizer_with_param[0][-1]))
+                        self.loss_list.append(loss_value.item())
+                        Logger.debug(f'last loss value: {self.loss_list[-1]}')
 
     def update_graph(self) -> Graph:
         """

model_compression_toolkit/gptq/pytorch/quantization_facade.py

@@ -31,6 +31,7 @@ from model_compression_toolkit.core.analyzer import analyzer_model_quantization
 from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
     MixedPrecisionQuantizationConfig
+from model_compression_toolkit.metadata import get_versions_dict
 
 LR_DEFAULT = 1e-4
 LR_REST_DEFAULT = 1e-4
@@ -47,6 +48,7 @@ if FOUND_TORCH:
     from torch.nn import Module
     from torch.optim import Adam, Optimizer
     from model_compression_toolkit import get_target_platform_capabilities
+    from mct_quantizers.pytorch.metadata import add_metadata
     DEFAULT_PYTORCH_TPC = get_target_platform_capabilities(PYTORCH, DEFAULT_TP_MODEL)
 
     def get_pytorch_gptq_config(n_epochs: int,
@@ -202,7 +204,10 @@ if FOUND_TORCH:
                                         fw_impl,
                                         DEFAULT_PYTORCH_INFO)
 
-        return get_exportable_pytorch_model(graph_gptq)
+        exportable_model, user_info = get_exportable_pytorch_model(graph_gptq)
+        if target_platform_capabilities.tp_model.add_metadata:
+            exportable_model = add_metadata(exportable_model, get_versions_dict(target_platform_capabilities))
+        return exportable_model, user_info
 
 
 else:

model_compression_toolkit/gptq/pytorch/quantizer/regularization_factory.py

@@ -12,6 +12,7 @@
 # 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
@@ -35,7 +36,7 @@ def get_regularization(gptq_config: GradientPTQConfig, representative_data_gen:
     if gptq_config.rounding_type == RoundingType.SoftQuantizer:
         # dry run on the representative dataset to count number of batches
         num_batches = 0
-        for _ in representative_data_gen():
+        for _ in tqdm(representative_data_gen(), "GPTQ initialization"):
             num_batches += 1
 
         return SoftQuantizerRegularization(total_gradient_steps=num_batches * gptq_config.n_epochs)