mct-nightly 2.1.0.20240725.446__py3-none-any.whl → 2.1.0.20240727.431__py3-none-any.whl

{mct_nightly-2.1.0.20240725.446.dist-info → mct_nightly-2.1.0.20240727.431.dist-info}/METADATA

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

{mct_nightly-2.1.0.20240725.446.dist-info → mct_nightly-2.1.0.20240727.431.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=vR8xewBS60bSjPLOXnFBfDT7srZ_io2XFWfqLG-lvkw,1573
+model_compression_toolkit/__init__.py,sha256=dNSTIwKIETFrRFbKVtmCp7sMOFv7eHAfRFVw0joUkik,1573
 model_compression_toolkit/constants.py,sha256=9pVleMwnhlM4QwIL2HcEq42I1uF4rlSw63RUjkxOF4w,3923
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
@@ -210,13 +210,13 @@ model_compression_toolkit/core/keras/statistics_correction/__init__.py,sha256=9H
 model_compression_toolkit/core/keras/statistics_correction/apply_second_moment_correction.py,sha256=XNCtT9klMcsO1v5KA3MmCq_WgXOIT5QSzbfTOa9T-04,3060
 model_compression_toolkit/core/keras/visualization/__init__.py,sha256=mjbqLD-KcG3eNeCYpu1GBS7VclGVOQ63x2p6mAAuba4,698
 model_compression_toolkit/core/pytorch/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
-model_compression_toolkit/core/pytorch/constants.py,sha256=AguUnAsNlj41gwuKIP_7nos3FcJHsIAjewLXSQdrDQM,2624
+model_compression_toolkit/core/pytorch/constants.py,sha256=YwD_joIF0vK8UG2vW1NVvg36pCNWA0vHOXjAgy_XWn0,2794
 model_compression_toolkit/core/pytorch/default_framework_info.py,sha256=r1XyzUFvrjGcJHQM5ETLsMZIG2yHCr9HMjqf0ti9inw,4175
 model_compression_toolkit/core/pytorch/pytorch_device_config.py,sha256=S25cuw10AW3SEN_fRAGRcG_I3wdvvQx1ehSJzPnn-UI,4404
 model_compression_toolkit/core/pytorch/pytorch_implementation.py,sha256=XL_RZcfnb_ZY2jdCjOxxz7SbRBzMokbOWsTuYOSjyRU,27569
 model_compression_toolkit/core/pytorch/pytorch_node_prior_info.py,sha256=2LDQ7qupglHQ7o1Am7LWdfYVacfQnl-aW2N6l9det1w,3264
 model_compression_toolkit/core/pytorch/resource_utilization_data_facade.py,sha256=E6ifk1HdO60k4IRH2EFBzAYWtwUlrGqJoQ66nknpHoQ,4983
-model_compression_toolkit/core/pytorch/utils.py,sha256=OT_mrNEJqPgWLdtQuivKMQVjtJY49cmoIVvbRhANl1w,3004
+model_compression_toolkit/core/pytorch/utils.py,sha256=GE7T8q93I5C4As0iOias_dk9HpOvXM1N6---dJlyD60,3863
 model_compression_toolkit/core/pytorch/back2framework/__init__.py,sha256=H_WixgN0elVWf3exgGYsi58imPoYDj5eYPeh6x4yfug,813
 model_compression_toolkit/core/pytorch/back2framework/factory_model_builder.py,sha256=DwNO8WO3JiMawKGKDhlrwCoCjMSBIw5BMbsFFF7eDS4,2279
 model_compression_toolkit/core/pytorch/back2framework/float_model_builder.py,sha256=tLrlUyYhxVKVjkad1ZAtbRra0HedB3iVfIkZ_dYnQ-4,3419
@@ -269,36 +269,40 @@ model_compression_toolkit/core/pytorch/statistics_correction/apply_second_moment
 model_compression_toolkit/data_generation/__init__.py,sha256=S8pRUqlRvpM5AFHpFWs3zb0H0rtY5nUwmeCQij01oi4,1507
 model_compression_toolkit/data_generation/common/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/data_generation/common/constants.py,sha256=21e3ZX9WVYojexG2acTgklrBk8ZO9DjJnKpP4KHZC44,1018
-model_compression_toolkit/data_generation/common/data_generation.py,sha256=fccGG6cTMScZwjnJDQKMugOLdgm9dKg5rRfcBD4EFYQ,6415
-model_compression_toolkit/data_generation/common/data_generation_config.py,sha256=ynyNaT2x2d23bYSrO2sRItM2ZsjGD0K0fM71FlibiJQ,4564
-model_compression_toolkit/data_generation/common/enums.py,sha256=OGnvtEGFbP5l4V3-1l32zzVQwTb1vGJhTVF0kOkYZK4,3584
-model_compression_toolkit/data_generation/common/image_pipeline.py,sha256=WwyeoIvgmcxKnuOX-_Hl_0APET4M26f5x-prhUB3qvU,2149
-model_compression_toolkit/data_generation/common/model_info_exctractors.py,sha256=kz3w4h4fO4R2N5IgLvSkqDUwjhH4S-I3n3_pK2hQ_uo,6200
-model_compression_toolkit/data_generation/common/optimization_utils.py,sha256=Q_yG8T8HQxfTKjVmN7bLm4M4y-5WrMeGQ_h5RnurSmg,19761
+model_compression_toolkit/data_generation/common/data_generation.py,sha256=W8PeOcL1fBVB1WgXSCrEw-G7AWa6MNzjTqcFbmMhrGE,6687
+model_compression_toolkit/data_generation/common/data_generation_config.py,sha256=yKqSDJGdbnc9HEmg94sPqMSXGR2OmAzt5X5MQcy_YX8,4473
+model_compression_toolkit/data_generation/common/enums.py,sha256=V5qAaqMg2WFhsrJ11rTDcRWBhbsxhEHt3uwRq6cesNo,4249
+model_compression_toolkit/data_generation/common/image_pipeline.py,sha256=PfunQMxYm6KqJUEUVYhtY7-JTq4J-XTyLc1HOalP15s,4761
+model_compression_toolkit/data_generation/common/model_info_exctractors.py,sha256=CqruljgQ564SMRQtxgYYDWKM7HYDz18MCShNgrRYQKg,5933
+model_compression_toolkit/data_generation/common/optimization_utils.py,sha256=aEDSclZ2TvIIqN1x9CLf8MBe2GA3m1aEXtbd5Sgcd8k,19528
 model_compression_toolkit/data_generation/keras/__init__.py,sha256=lNJ29DYxaLUPDstRDA1PGI5r9Fulq_hvrZMlhst1Z5g,697
-model_compression_toolkit/data_generation/keras/constants.py,sha256=uy3eU24ykygIrjIvwOMj3j5euBeN2PwWiEFPOkJJ7ss,1088
-model_compression_toolkit/data_generation/keras/image_pipeline.py,sha256=_Qezq67huKmmNsxdFBBrTY-VaGR-paFzDH80dDuRnug,7623
-model_compression_toolkit/data_generation/keras/keras_data_generation.py,sha256=lWNxHIews7fREVGT99b2EPV6VdpuLjwDINLApatlPjo,21539
-model_compression_toolkit/data_generation/keras/model_info_exctractors.py,sha256=d_SOttfnOk97_nJ7Vjv5p-8tcQqAXJ_hzc1khZGU92s,8610
-model_compression_toolkit/data_generation/keras/optimization_utils.py,sha256=uQAJpJPpnLDTTLDQGyTS0ZYp2T38TTZLOOElcJPBKHA,21146
+model_compression_toolkit/data_generation/keras/constants.py,sha256=sxhhGHC-INBs1nVXhyokbFi9ob4jPkSRviuc83JRsgQ,1152
+model_compression_toolkit/data_generation/keras/image_operations.py,sha256=OtJ5Yz8BZVOnGqyTHwlseRe4EmoLDYxz3bblGtw6HnY,6233
+model_compression_toolkit/data_generation/keras/image_pipeline.py,sha256=E-HVverorhq33xzteuwUPtOrGDIYoEEs4fZJgiqOAzQ,7043
+model_compression_toolkit/data_generation/keras/keras_data_generation.py,sha256=IMnmUn7fUsMcJ980FZWuX36iUYXAEYxdYk8oXwz-Xd8,21207
+model_compression_toolkit/data_generation/keras/model_info_exctractors.py,sha256=1E5xbn0P3py4EYjdpPD9JwGr4jlc3qe1ml1py0t40b8,8026
+model_compression_toolkit/data_generation/keras/optimization_utils.py,sha256=cHv2tl-_9_D14mWqzNYtKFY8q7sJfW_V__dpZqzRvIo,20546
 model_compression_toolkit/data_generation/keras/optimization_functions/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/data_generation/keras/optimization_functions/batchnorm_alignment_functions.py,sha256=f5M7KoISGnb6S6zR7SyQ9dYmQctW9iYRi0Bda1BLq70,1983
-model_compression_toolkit/data_generation/keras/optimization_functions/bn_layer_weighting_functions.py,sha256=T6P7VEhRRcJAvpPhMOp3izo8Gn13_nMb7GNNm1w0WhA,2699
+model_compression_toolkit/data_generation/keras/optimization_functions/bn_layer_weighting_functions.py,sha256=xQWTeP-Im6xEUupF-VEjZq-UsRNzpoW0LuMHFR2cX9Q,3390
 model_compression_toolkit/data_generation/keras/optimization_functions/image_initilization.py,sha256=sjSPLLFLjJ6d0DDSaxnCE0ydIT1zhL8H73QTXEuUfgw,4119
-model_compression_toolkit/data_generation/keras/optimization_functions/output_loss_functions.py,sha256=bVb6pnVzGPzC7_Hjj1X3WUUNSfndYUG8_PL_-nQo2UU,5325
-model_compression_toolkit/data_generation/keras/optimization_functions/scheduler_step_functions.py,sha256=1HK5-Z0WSk_L62RB2RQGlrwKpCQyoR2TKI5qjXod6cc,4913
+model_compression_toolkit/data_generation/keras/optimization_functions/lr_scheduler.py,sha256=xCc7GlmW-jpf27P8mI2APpAL8LC6zmD1BMbN7Q6wzEE,8647
+model_compression_toolkit/data_generation/keras/optimization_functions/output_loss_functions.py,sha256=vr_H1dbFINS7LBX_SfW59g0C8ie9grAyOIpCKuPoI1w,6384
+model_compression_toolkit/data_generation/keras/optimization_functions/scheduler_step_functions.py,sha256=9RhNWtw_cdDlGqEGEdn1JWwvfA8V-Z6ioZn1ppdHFmA,1695
 model_compression_toolkit/data_generation/pytorch/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
-model_compression_toolkit/data_generation/pytorch/constants.py,sha256=QWyreMImcfKzW1RBqf-4yl2kAr4sAABNNgQmPUkqPEo,1179
-model_compression_toolkit/data_generation/pytorch/image_pipeline.py,sha256=6g7OpOuO3cU4TIuelaRjBKpCPgiMbe1a3iy9bZtdZUo,6617
-model_compression_toolkit/data_generation/pytorch/model_info_exctractors.py,sha256=sO9tA03nIaeYnzOL4Egec5sVcSGU8H8k9-nNjhaLEbk,9690
-model_compression_toolkit/data_generation/pytorch/optimization_utils.py,sha256=AjYsO-lm06JOUMoKkS6VbyF4O_l_ffWXrgamqJm1ofE,19085
-model_compression_toolkit/data_generation/pytorch/pytorch_data_generation.py,sha256=Jymevochf1A6sz_bShQiJVj5IUtxKbfW5s5Bh7EZhUo,21238
+model_compression_toolkit/data_generation/pytorch/constants.py,sha256=ZiyweWj2Bnk6duhcV4zowsPvqLdON-AlLhkAuLmCqxg,1256
+model_compression_toolkit/data_generation/pytorch/image_operations.py,sha256=KUQKOj5G4UPGX9f9PSiLRlBo4e3rRRPec88wkozNgqw,3900
+model_compression_toolkit/data_generation/pytorch/image_pipeline.py,sha256=dcQr-67u9-ggGuS39YAvR7z-Y0NOdJintcVQ5vy1bM8,7478
+model_compression_toolkit/data_generation/pytorch/model_info_exctractors.py,sha256=y6vMed6lQQj67-BXZKrAcWUNTkH8YjiUhknOV4wSpRA,9399
+model_compression_toolkit/data_generation/pytorch/optimization_utils.py,sha256=vRMeUEdInPuJisiO-SKo_9miWZV90sz8GCg5MY0AqiU,18098
+model_compression_toolkit/data_generation/pytorch/pytorch_data_generation.py,sha256=OjdAG0uGdbN0ATMrkRskhEttkUgSXN8KCVd8JXKiwxk,21620
 model_compression_toolkit/data_generation/pytorch/optimization_functions/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/data_generation/pytorch/optimization_functions/batchnorm_alignment_functions.py,sha256=dMc4zz9XfYfAT4Cxns57VgvGZWPAMfaGlWLFyCyl8TA,1968
-model_compression_toolkit/data_generation/pytorch/optimization_functions/bn_layer_weighting_functions.py,sha256=i3ePEI8xDE3xZEtmzT5lCkLn9wpObUi_OgqnVDf7nj8,2597
-model_compression_toolkit/data_generation/pytorch/optimization_functions/image_initilization.py,sha256=gQDZvuT2yt3V1K5TPvBwAdvdTU-ztsQgGAL7ErKm1PQ,4719
-model_compression_toolkit/data_generation/pytorch/optimization_functions/output_loss_functions.py,sha256=75_Dl7clEa_SDYRVF3Ub27urYMar8KydZyTUEvZT3ks,5015
-model_compression_toolkit/data_generation/pytorch/optimization_functions/scheduler_step_functions.py,sha256=bxbNlQY7AKX8lEJaJReEn0lGBGb6e4PiDEqyW__nmig,2587
+model_compression_toolkit/data_generation/pytorch/optimization_functions/bn_layer_weighting_functions.py,sha256=We0fVMQ4oU7Y0IWQ8fKy8KpqkIiLyKoQeF9XKAQ6TH0,3317
+model_compression_toolkit/data_generation/pytorch/optimization_functions/image_initilization.py,sha256=hhWSZ7w45dE5SQ6jM27cBkCSJObWkALs_RpD6afPi68,4753
+model_compression_toolkit/data_generation/pytorch/optimization_functions/lr_scheduler.py,sha256=NydGxFIclmrfU3HWYUrRbprg4hPt470QP6MTOMLEhRs,9172
+model_compression_toolkit/data_generation/pytorch/optimization_functions/output_loss_functions.py,sha256=PRVmn8o2hTdwTdbd2ezf__LNbFvcgiVO0c25dsyg3Tg,6549
+model_compression_toolkit/data_generation/pytorch/optimization_functions/scheduler_step_functions.py,sha256=zMjY2y4FSHonuY5hddbMTb8qAQtLtohYF7q1wuruDDs,3267
 model_compression_toolkit/exporter/__init__.py,sha256=Eg3c4EAjW3g6h13A-Utgf9ncHrTMRHAoySNDQGPDZ4E,1301
 model_compression_toolkit/exporter/model_exporter/__init__.py,sha256=9HIBmj8ROdCA-yvkpA8EcN6RHJe_2vEpLLW_gxOJtak,698
 model_compression_toolkit/exporter/model_exporter/fw_agonstic/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
@@ -517,8 +521,8 @@ model_compression_toolkit/xquant/pytorch/model_analyzer.py,sha256=b93o800yVB3Z-i
 model_compression_toolkit/xquant/pytorch/pytorch_report_utils.py,sha256=yrZNVRm2IRU7r7R-hjS2lOQ6wvEEvbeunvf2jKoWjXk,3277
 model_compression_toolkit/xquant/pytorch/similarity_functions.py,sha256=CERxq5K8rqaiE-DlwhZBTUd9x69dtYJlkHOPLB54vm8,2354
 model_compression_toolkit/xquant/pytorch/tensorboard_utils.py,sha256=eyMoXt5o5EnMr6d-rpCwQdX5mAiYiymvbgKv4tf7-a0,4576
-mct_nightly-2.1.0.20240725.446.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.1.0.20240725.446.dist-info/METADATA,sha256=m0FhYyaViVkXqL-az7oYXCs-wZpStJIJO_Tbf0wISUM,19719
-mct_nightly-2.1.0.20240725.446.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-mct_nightly-2.1.0.20240725.446.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.1.0.20240725.446.dist-info/RECORD,,
+mct_nightly-2.1.0.20240727.431.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.1.0.20240727.431.dist-info/METADATA,sha256=g2Z7XHdZo_A9_vL0R9IJv5kcLIsZg_ONvFBnuOFqFkk,19719
+mct_nightly-2.1.0.20240727.431.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+mct_nightly-2.1.0.20240727.431.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.1.0.20240727.431.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.1.0.20240725.000446"
+__version__ = "2.1.0.20240727.000431"

model_compression_toolkit/core/pytorch/constants.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
+import torch
 
 # # Layer type constants:
 PLACEHOLDER = 'placeholder'
@@ -94,3 +94,8 @@ BIAS_V = 'bias_v'
 # # Batch size value for 'reshape' and 'view' operators,
 # # the value is -1 so the batch size is inferred from the length of the array and remaining dimensions.
 BATCH_DIM_VALUE = -1
+
+
+# The maximum and minimum representable values for float16
+MAX_FLOAT16 = torch.finfo(torch.float16).max - 1
+MIN_FLOAT16 = torch.finfo(torch.float16).min - 1

model_compression_toolkit/core/pytorch/utils.py

@@ -13,8 +13,11 @@
 # limitations under the License.
 # ==============================================================================
 import torch
+from torch import Tensor
 import numpy as np
 from typing import Union
+
+from model_compression_toolkit.core.pytorch.constants import MAX_FLOAT16, MIN_FLOAT16
 from model_compression_toolkit.core.pytorch.pytorch_device_config import get_working_device
 from model_compression_toolkit.logger import Logger
 
@@ -83,3 +86,27 @@ def torch_tensor_to_numpy(tensor: Union[torch.Tensor, list, tuple]) -> Union[np.
         return tensor.cpu().detach().contiguous().numpy()
     else:
         Logger.critical(f'Unsupported type for conversion to Numpy array: {type(tensor)}.')
+
+
+def clip_inf_values_float16(tensor: Tensor) -> Tensor:
+    """
+    Clips +inf and -inf values in a float16 tensor to the maximum and minimum representable values.
+
+    Parameters:
+    tensor (Tensor): Input PyTorch tensor of dtype float16.
+
+    Returns:
+    Tensor: A tensor with +inf values replaced by the maximum float16 value,
+            and -inf values replaced by the minimum float16 value.
+    """
+    # Check if the tensor is of dtype float16
+    if tensor.dtype != torch.float16:
+        return tensor
+
+    # Create a mask for inf values (both positive and negative)
+    inf_mask = torch.isinf(tensor)
+
+    # Replace inf values with max float16 value
+    tensor[inf_mask] = MAX_FLOAT16 * torch.sign(tensor[inf_mask])
+
+    return tensor

model_compression_toolkit/data_generation/common/data_generation.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 # Import required modules and classes
-from typing import Any, Tuple, Dict, Callable, List
+from typing import Any, Tuple, Dict, Callable, List, Union
 
 from model_compression_toolkit.data_generation.common.data_generation_config import DataGenerationConfig
 from model_compression_toolkit.data_generation.common.enums import ImagePipelineType, ImageNormalizationType, \
@@ -24,7 +24,7 @@ from model_compression_toolkit.logger import Logger
 
 def get_data_generation_classes(
         data_generation_config: DataGenerationConfig,
-        output_image_size: Tuple,
+        output_image_size: Union[int, Tuple[int, int]],
         n_images: int,
         image_pipeline_dict: Dict,
         image_normalization_dict: Dict,
@@ -38,7 +38,7 @@ def get_data_generation_classes(
 
     Args:
         data_generation_config (DataGenerationConfig): Configuration for data generation.
-        output_image_size (Tuple): The desired output image size.
+        output_image_size (Union[int, Tuple[int, int]]): The desired output image size.
         n_images (int): The number of random samples.
         image_pipeline_dict (Dict): Dictionary mapping ImagePipelineType to corresponding image pipeline classes.
         image_normalization_dict (Dict): Dictionary mapping ImageNormalizationType to corresponding
@@ -56,26 +56,28 @@ def get_data_generation_classes(
         output_loss_fn (Callable): Function to compute output loss.
         init_dataset (Any): The initial dataset used for image generation.
     """
+    # Get the normalization values corresponding to the specified type
+    normalization = image_normalization_dict.get(data_generation_config.image_normalization_type)
+
+    # Check if the image normalization type is valid
+    if normalization is None:
+        Logger.critical(
+            f'Invalid image_normalization_type {data_generation_config.image_normalization_type}. '
+            f'Please select one from {ImageNormalizationType.get_values()}.') # pragma: no cover
+
     # Get the image pipeline class corresponding to the specified type
     image_pipeline = (
         image_pipeline_dict.get(data_generation_config.image_pipeline_type)(
             output_image_size=output_image_size,
-            extra_pixels=data_generation_config.extra_pixels))
+            extra_pixels=data_generation_config.extra_pixels,
+            image_clipping=data_generation_config.image_clipping,
+            normalization=normalization))
 
     # Check if the image pipeline type is valid
     if image_pipeline is None:
         Logger.critical(
             f'Invalid image_pipeline_type {data_generation_config.image_pipeline_type}. '
-            f'Please select one from {ImagePipelineType.get_values()}.')
-
-    # Get the normalization values corresponding to the specified type
-    normalization = image_normalization_dict.get(data_generation_config.image_normalization_type)
-
-    # Check if the image normalization type is valid
-    if normalization is None:
-        Logger.critical(
-            f'Invalid image_normalization_type {data_generation_config.image_normalization_type}. '
-            f'Please select one from {ImageNormalizationType.get_values()}.')
+            f'Please select one from {ImagePipelineType.get_values()}.') # pragma: no cover
 
     # Get the layer weighting function corresponding to the specified type
     bn_layer_weighting_fn = bn_layer_weighting_function_dict.get(data_generation_config.layer_weighting_type)
@@ -83,7 +85,7 @@ def get_data_generation_classes(
     if bn_layer_weighting_fn is None:
         Logger.critical(
             f'Invalid layer_weighting_type {data_generation_config.layer_weighting_type}. '
-            f'Please select one from {BNLayerWeightingType.get_values()}.')
+            f'Please select one from {BNLayerWeightingType.get_values()}.') # pragma: no cover
 
     # Get the image initialization function corresponding to the specified type
     image_initialization_fn = image_initialization_function_dict.get(data_generation_config.data_init_type)
@@ -92,7 +94,7 @@ def get_data_generation_classes(
     if image_initialization_fn is None:
         Logger.critical(
             f'Invalid data_init_type {data_generation_config.data_init_type}. '
-            f'Please select one from {DataInitType.get_values()}.')
+            f'Please select one from {DataInitType.get_values()}.') # pragma: no cover
 
     # Get the BatchNorm alignment loss function corresponding to the specified type
     bn_alignment_loss_fn = bn_alignment_loss_function_dict.get(data_generation_config.bn_alignment_loss_type)
@@ -101,7 +103,7 @@ def get_data_generation_classes(
     if bn_alignment_loss_fn is None:
         Logger.critical(
             f'Invalid bn_alignment_loss_type {data_generation_config.bn_alignment_loss_type}. '
-            f'Please select one from {BatchNormAlignemntLossType.get_values()}.')
+            f'Please select one from {BatchNormAlignemntLossType.get_values()}.') # pragma: no cover
 
     # Get the output loss function corresponding to the specified type
     output_loss_fn = output_loss_function_dict.get(data_generation_config.output_loss_type)
@@ -110,7 +112,7 @@ def get_data_generation_classes(
     if output_loss_fn is None:
         Logger.critical(
             f'Invalid output_loss_type {data_generation_config.output_loss_type}. '
-            f'Please select one from {OutputLossType.get_values()}.')
+            f'Please select one from {OutputLossType.get_values()}.') # pragma: no cover
 
     # Initialize the dataset for data generation
     init_dataset = image_initialization_fn(

model_compression_toolkit/data_generation/common/data_generation_config.py

@@ -12,12 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Callable, Any, List
+from typing import Any, List, Tuple, Union
 
 from model_compression_toolkit.data_generation.common.enums import SchedulerType, BatchNormAlignemntLossType, \
     DataInitType, BNLayerWeightingType, ImageGranularity, ImagePipelineType, ImageNormalizationType, OutputLossType
 
-
 class DataGenerationConfig:
     """
     Configuration class for data generation.
@@ -36,11 +35,10 @@ class DataGenerationConfig:
                  layer_weighting_type: BNLayerWeightingType = None,
                  image_pipeline_type: ImagePipelineType = None,
                  image_normalization_type: ImageNormalizationType = None,
-                 extra_pixels: int = 0,
+                 extra_pixels: Union[int, Tuple[int, int]] = 0,
                  bn_layer_types: List = [],
                  last_layer_types: List = [],
-                 clip_images: bool = True,
-                 reflection: bool = True,
+                 image_clipping: bool = True,
                  ):
         """
         Initialize the DataGenerationConfig.
@@ -59,17 +57,15 @@ class DataGenerationConfig:
             layer_weighting_type (BNLayerWeightingType): Type of layer weighting. Defaults to None.
             image_pipeline_type (ImagePipelineType): Type of image pipeline. Defaults to None.
             image_normalization_type (ImageNormalizationType): Type of image normalization. Defaults to None.
-            extra_pixels (int): Extra pixels to add to the input image size. Defaults to 0.
+            extra_pixels (Union[int, Tuple[int, int]]): Extra pixels to add to the input image size. Defaults to 0.
             bn_layer_types (List): List of BatchNorm layer types. Defaults to [].
             last_layer_types (List): List of layer types. Defaults to [].
-            clip_images (bool): Flag to enable image clipping. Defaults to True.
-            reflection (bool): Flag to enable reflection. Defaults to True.
+            image_clipping (bool): Flag to enable image clipping. Defaults to True.
         """
         self.n_iter = n_iter
         self.optimizer = optimizer
         self.data_gen_batch_size = data_gen_batch_size
         self.initial_lr = initial_lr
-        self.output_loss_multiplier = output_loss_multiplier
         self.image_granularity = image_granularity
         self.scheduler_type = scheduler_type
         self.bn_alignment_loss_type = bn_alignment_loss_type
@@ -81,6 +77,7 @@ class DataGenerationConfig:
         self.layer_weighting_type = layer_weighting_type
         self.bn_layer_types = bn_layer_types
         self.last_layer_types = last_layer_types
-        self.clip_images = clip_images
-        self.reflection = reflection
+        self.image_clipping = image_clipping
+        self.output_loss_multiplier = output_loss_multiplier
+
 

model_compression_toolkit/data_generation/common/enums.py

@@ -16,7 +16,6 @@ from enum import Enum
 
 
 class EnumBaseClass(Enum):
-    @classmethod
     def get_values(cls):
         """
         Get the list of values corresponding to the enum members.
@@ -24,7 +23,23 @@ class EnumBaseClass(Enum):
         Returns:
             List of values.
         """
-        return [value.value for value in cls.__members__.values()]
+        return list(cls.__members__.values())
+
+    @classmethod
+    def get_enum_by_value(cls, target_value):
+        """
+        Function to get the key corresponding to a given enum value.
+
+        Parameters:
+        target_value: The value to find the key for.
+
+        Returns:
+        The key corresponding to the given value if found, else None.
+        """
+        for value in cls.__members__.values():
+            if value.value == target_value:
+                return value
+        return None
 
 
 class ImageGranularity(EnumBaseClass):
@@ -61,15 +76,12 @@ class ImagePipelineType(EnumBaseClass):
     """
     An enum for choosing the image pipeline type for image manipulation:
 
-    RANDOM_CROP - Crop the images.
-
-    RANDOM_CROP_FLIP - Crop and flip the images.
+    SMOOTHING_AND_AUGMENTATION - Apply a smoothing filter, then crop and flip the images.
 
     IDENTITY - Do not apply any manipulation (identity transformation).
 
     """
-    RANDOM_CROP = 'random_crop'
-    RANDOM_CROP_FLIP = 'random_crop_flip'
+    SMOOTHING_AND_AUGMENTATION = 'smoothing_and_augmentation'
     IDENTITY = 'identity'
 
 
@@ -118,16 +130,15 @@ class BatchNormAlignemntLossType(EnumBaseClass):
 class OutputLossType(EnumBaseClass):
     """
     An enum for choosing the output loss type:
-
     NONE - No output loss is applied.
-
-    MIN_MAX_DIFF - Use min-max difference as the output loss.
-
+    NEGATIVE_MIN_MAX_DIFF - Use the mean of the negative min-max difference as the output loss.
+    INVERSE_MIN_MAX_DIFF - Use mean of the 1/(min-max) difference as the output loss.
     REGULARIZED_MIN_MAX_DIFF - Use regularized min-max difference as the output loss.
 
     """
     NONE = 'none'
-    MIN_MAX_DIFF = 'min_max_diff'
+    NEGATIVE_MIN_MAX_DIFF = 'negative_min_max_diff'
+    INVERSE_MIN_MAX_DIFF = 'inverse_min_max_diff'
     REGULARIZED_MIN_MAX_DIFF = 'regularized_min_max_diff'
 
 
@@ -141,4 +152,5 @@ class SchedulerType(EnumBaseClass):
 
     """
     REDUCE_ON_PLATEAU = 'reduce_on_plateau'
+    REDUCE_ON_PLATEAU_WITH_RESET = 'reduce_on_plateau_with_reset'
     STEP = 'step'

model_compression_toolkit/data_generation/common/image_pipeline.py

@@ -13,31 +13,61 @@
 # limitations under the License.
 # ==============================================================================
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, Tuple, Union, Dict, List
+
+from model_compression_toolkit.data_generation import ImageNormalizationType
+from model_compression_toolkit.logger import Logger
 
 
 class BaseImagePipeline(ABC):
     def __init__(self,
-                 output_image_size: int,
-                 extra_pixels: int = 0):
+                 output_image_size: Union[int, Tuple[int, int]],
+                 extra_pixels: Union[int, Tuple[int, int]] = 0,
+                 image_clipping: bool = False,
+                 normalization: List[List[int]] = [[0, 0, 0], [1, 1, 1]]):
         """
         Base class for image pipeline.
 
         Args:
-            output_image_size (int): The desired output image size.
-            extra_pixels (int, optional): Extra pixels to add to the input image size. Defaults to 0.
+            output_image_size (Union[int, Tuple[int, int]]): The desired output image size.
+            extra_pixels (Union[int, Tuple[int, int]]): Extra pixels to add to the input image size. Defaults to 0.
+            image_clipping (bool): Whether to clip images during optimization.
+            normalization (List[List[float]]): The image normalization values for processing images during optimization.
         """
-        self.output_image_size = output_image_size
-        self.extra_pixels = extra_pixels
+        if isinstance(output_image_size, int):
+            self.output_image_size = (output_image_size, output_image_size)
+        elif isinstance(output_image_size, tuple) and len(output_image_size) == 1:
+            self.output_image_size = output_image_size + output_image_size # concatenate two tuples
+        elif isinstance(output_image_size, tuple) and len(output_image_size) == 2:
+            self.output_image_size = output_image_size
+        elif isinstance(output_image_size, tuple):
+            Logger.critical(f"'output_image_size' should a tuple of length 1 or 2. Got tuple of length {len(output_image_size)}") # pragma: no cover
+        else:
+            Logger.critical(f"'output_image_size' should be an int or tuple but type {type(output_image_size)} was received.") # pragma: no cover
+
+        if isinstance(extra_pixels, int):
+            self.extra_pixels = (extra_pixels, extra_pixels)
+        elif isinstance(extra_pixels, tuple) and len(extra_pixels) == 1:
+            self.extra_pixels = extra_pixels + extra_pixels # concatenate two tuples
+        elif isinstance(extra_pixels, tuple) and len(extra_pixels) == 2:
+            self.extra_pixels = extra_pixels
+        elif isinstance(extra_pixels, tuple):
+            Logger.critical(f"'extra_pixels' should a tuple of length 1 or 2. Got tuple of length {len(extra_pixels)}") # pragma: no cover
+        else:
+            Logger.critical(f"'extra_pixels' should be an int or tuple but type {type(extra_pixels)} was received.") # pragma: no cover
+
+        self.image_clipping = image_clipping
+        self.normalization = normalization
+
     @abstractmethod
-    def get_image_input_size(self) -> int:
+    def get_image_input_size(self) -> Tuple[int, int]:
         """
         Get the size of the input image for the image pipeline.
 
         Returns:
-            int: The input image size.
+            Tuple[int, int]: The input image size.
         """
-        raise NotImplemented
+        raise NotImplemented # pragma: no cover
 
     @abstractmethod
     def image_input_manipulation(self,
@@ -51,7 +81,7 @@ class BaseImagePipeline(ABC):
         Returns:
             Any: Manipulated images.
         """
-        raise NotImplemented
+        raise NotImplemented # pragma: no cover
 
     @abstractmethod
     def image_output_finalize(self,
@@ -65,4 +95,12 @@ class BaseImagePipeline(ABC):
         Returns:
             Any: Finalized images.
         """
-        raise NotImplemented
+        raise NotImplemented # pragma: no cover
+
+
+# Dictionary mapping ImageNormalizationType to corresponding normalization values
+image_normalization_dict: Dict[ImageNormalizationType, List[List[float]]] = {
+    ImageNormalizationType.TORCHVISION: [[0.485 * 255, 0.456 * 255, 0.406 * 255], [0.229 * 255, 0.224 * 255, 0.225 * 255]],
+    ImageNormalizationType.KERAS_APPLICATIONS: [[127.5, 127.5, 127.5], [127.5, 127.5, 127.5]],
+    ImageNormalizationType.NO_NORMALIZATION: [[0, 0, 0], [1, 1, 1]]
+}

model_compression_toolkit/data_generation/common/model_info_exctractors.py

@@ -157,14 +157,6 @@ class ActivationExtractor:
         """
         raise NotImplemented # pragma: no cover
 
-    def get_num_extractor_layers(self) -> int:
-        """
-        Get the number of layers for which to extract input activations.
-
-        Returns:
-            int: Number of layers for which to extract input activations.
-        """
-        return self.num_layers
 
     @abstractmethod
     def get_extractor_layer_names(self) -> List:

model_compression_toolkit/data_generation/common/optimization_utils.py

@@ -48,8 +48,6 @@ class ImagesOptimizationHandler:
                  scheduler: Any,
                  normalization_mean: List[float],
                  normalization_std: List[float],
-                 clip_images: bool,
-                 reflection: bool,
                  initial_lr: float,
                  eps: float = 1e-6):
         """
@@ -67,8 +65,6 @@ class ImagesOptimizationHandler:
             scheduler (Any): The scheduler responsible for adjusting the learning rate of the optimizer over time.
             normalization_mean (List[float]): Mean values used for image normalization.
             normalization_std (List[float]): Standard deviation values used for image normalization.
-            clip_images (bool): Flag indicating whether to clip generated images during optimization.
-            reflection (bool): Flag indicating whether to use reflection during image generation.
             initial_lr (float): The initial learning rate used by the optimizer.
             eps (float, optional): A small value added for numerical stability. Defaults to 1e-6.
 
@@ -79,8 +75,6 @@ class ImagesOptimizationHandler:
         self.scheduler = scheduler
         self.scheduler_step_fn = scheduler_step_fn
         self.image_granularity = image_granularity
-        self.clip_images = clip_images
-        self.reflection = reflection
         self.eps = eps
         self.targets = []
         self.initial_lr = initial_lr
@@ -209,9 +203,11 @@ class ImagesOptimizationHandler:
                 imgs_layer_mean, imgs_layer_second_moment, imgs_layer_std = self.all_imgs_stats_holder.get_stats(
                     batch_index, layer_name)
 
-            # Accumulate the batchnorm alignment weighted by the layer weight
-            total_bn_loss += bn_layer_weight * bn_alignment_loss_fn(bn_layer_mean, imgs_layer_mean, bn_layer_std,
-                                                                    imgs_layer_std)
+            if imgs_layer_mean is not None and imgs_layer_std is not None:
+                bn_alignment_loss = bn_alignment_loss_fn(bn_layer_mean, imgs_layer_mean, bn_layer_std,
+                                     imgs_layer_std)
+                # Accumulate the batchnorm alignment weighted by the layer weight
+                total_bn_loss += bn_layer_weight * bn_alignment_loss
 
         return total_bn_loss
 
@@ -418,7 +414,7 @@ class BatchStatsHolder:
         Returns:
             Any: the mean for the specified layer.
         """
-        return self.bn_mean[bn_layer_name]
+        return self.bn_mean.get(bn_layer_name)
 
     def get_second_moment(self, bn_layer_name: str) -> Any:
         """
@@ -430,7 +426,7 @@ class BatchStatsHolder:
         Returns:
             Any: the second moment for the specified layer.
         """
-        return self.bn_second_moment[bn_layer_name]
+        return self.bn_second_moment.get(bn_layer_name)
 
     def get_var(self, bn_layer_name: str) -> Any:
         """

model_compression_toolkit/data_generation/keras/constants.py

@@ -18,8 +18,11 @@ BATCH_AXIS, H_AXIS, W_AXIS, CHANNEL_AXIS = 0, 1, 2, 3
 # Default initial learning rate constant for Keras.
 DEFAULT_KERAS_INITIAL_LR = 1
 
-# Default output loss multiplier for Keras.
-DEFAULT_KERAS_OUTPUT_LOSS_MULTIPLIER = 0.001
+# Default extra pixels for image padding.
+DEFAULT_KERAS_EXTRA_PIXELS = 32
+
+# Default output loss multiplier.
+DEFAULT_KERAS_OUTPUT_LOSS_MULTIPLIER = 1e-3
 
 # Minimum value for image pixel intensity.
 IMAGE_MIN_VAL = 0.0