mct-nightly 2.1.0.20240724.437__tar.gz → 2.1.0.20240726.430__tar.gz

{mct-nightly-2.1.0.20240724.437 → mct-nightly-2.1.0.20240726.430}/PKG-INFO

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

{mct-nightly-2.1.0.20240724.437 → mct-nightly-2.1.0.20240726.430}/mct_nightly.egg-info/PKG-INFO

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

{mct-nightly-2.1.0.20240724.437 → mct-nightly-2.1.0.20240726.430}/mct_nightly.egg-info/SOURCES.txt

@@ -286,6 +286,7 @@ model_compression_toolkit/data_generation/common/model_info_exctractors.py
 model_compression_toolkit/data_generation/common/optimization_utils.py
 model_compression_toolkit/data_generation/keras/__init__.py
 model_compression_toolkit/data_generation/keras/constants.py
+model_compression_toolkit/data_generation/keras/image_operations.py
 model_compression_toolkit/data_generation/keras/image_pipeline.py
 model_compression_toolkit/data_generation/keras/keras_data_generation.py
 model_compression_toolkit/data_generation/keras/model_info_exctractors.py
@@ -294,10 +295,12 @@ model_compression_toolkit/data_generation/keras/optimization_functions/__init__.
 model_compression_toolkit/data_generation/keras/optimization_functions/batchnorm_alignment_functions.py
 model_compression_toolkit/data_generation/keras/optimization_functions/bn_layer_weighting_functions.py
 model_compression_toolkit/data_generation/keras/optimization_functions/image_initilization.py
+model_compression_toolkit/data_generation/keras/optimization_functions/lr_scheduler.py
 model_compression_toolkit/data_generation/keras/optimization_functions/output_loss_functions.py
 model_compression_toolkit/data_generation/keras/optimization_functions/scheduler_step_functions.py
 model_compression_toolkit/data_generation/pytorch/__init__.py
 model_compression_toolkit/data_generation/pytorch/constants.py
+model_compression_toolkit/data_generation/pytorch/image_operations.py
 model_compression_toolkit/data_generation/pytorch/image_pipeline.py
 model_compression_toolkit/data_generation/pytorch/model_info_exctractors.py
 model_compression_toolkit/data_generation/pytorch/optimization_utils.py
@@ -306,6 +309,7 @@ model_compression_toolkit/data_generation/pytorch/optimization_functions/__init_
 model_compression_toolkit/data_generation/pytorch/optimization_functions/batchnorm_alignment_functions.py
 model_compression_toolkit/data_generation/pytorch/optimization_functions/bn_layer_weighting_functions.py
 model_compression_toolkit/data_generation/pytorch/optimization_functions/image_initilization.py
+model_compression_toolkit/data_generation/pytorch/optimization_functions/lr_scheduler.py
 model_compression_toolkit/data_generation/pytorch/optimization_functions/output_loss_functions.py
 model_compression_toolkit/data_generation/pytorch/optimization_functions/scheduler_step_functions.py
 model_compression_toolkit/exporter/__init__.py

{mct-nightly-2.1.0.20240724.437 → mct-nightly-2.1.0.20240726.430}/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.20240724.000437"
+__version__ = "2.1.0.20240726.000430"

{mct-nightly-2.1.0.20240724.437 → mct-nightly-2.1.0.20240726.430}/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

{mct-nightly-2.1.0.20240724.437 → mct-nightly-2.1.0.20240726.430}/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

{mct-nightly-2.1.0.20240724.437 → mct-nightly-2.1.0.20240726.430}/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(

{mct-nightly-2.1.0.20240724.437 → mct-nightly-2.1.0.20240726.430}/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
+
 

{mct-nightly-2.1.0.20240724.437 → mct-nightly-2.1.0.20240726.430}/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'

mct-nightly-2.1.0.20240726.430/model_compression_toolkit/data_generation/common/image_pipeline.py

@@ -0,0 +1,106 @@
+# 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 abc import ABC, abstractmethod
+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: 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 (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.
+        """
+        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) -> Tuple[int, int]:
+        """
+        Get the size of the input image for the image pipeline.
+
+        Returns:
+            Tuple[int, int]: The input image size.
+        """
+        raise NotImplemented # pragma: no cover
+
+    @abstractmethod
+    def image_input_manipulation(self,
+                                 images: Any) -> Any:
+        """
+        Perform image input manipulation in the image pipeline.
+
+        Args:
+            images (Any): Input images.
+
+        Returns:
+            Any: Manipulated images.
+        """
+        raise NotImplemented # pragma: no cover
+
+    @abstractmethod
+    def image_output_finalize(self,
+                              images: Any) -> Any:
+        """
+        Perform finalization of output images in the image pipeline.
+
+        Args:
+            images (Any): Output images.
+
+        Returns:
+            Any: Finalized images.
+        """
+        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]]
+}

{mct-nightly-2.1.0.20240724.437 → mct-nightly-2.1.0.20240726.430}/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:

{mct-nightly-2.1.0.20240724.437 → mct-nightly-2.1.0.20240726.430}/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:
         """

{mct-nightly-2.1.0.20240724.437 → mct-nightly-2.1.0.20240726.430}/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