mct-nightly 2.1.0.20240724.437__py3-none-any.whl → 2.1.0.20240726.430__py3-none-any.whl

model_compression_toolkit/data_generation/keras/image_operations.py

@@ -0,0 +1,189 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import numpy as np
+from typing import Tuple
+
+import tensorflow as tf
+
+def random_crop(image: tf.Tensor,
+                height_crop: int,
+                width_crop: int) -> tf.Tensor:
+    """
+    Randomly crop an image to the specified size.
+
+    Args:
+        image (tf.Tensor): Input image tensor.
+        height_crop (int): Size of the crop in the height axis.
+        width_crop (int): Size of the crop in the width axis.
+
+    Returns:
+        tf.Tensor: Cropped image tensor.
+    """
+    cropped_image = tf.image.random_crop(image,
+                                         size=(tf.shape(image)[0],
+                                               height_crop,
+                                               width_crop,
+                                               tf.shape(image)[-1]))
+    return cropped_image
+
+
+def center_crop(image: tf.Tensor,
+                height_crop: int,
+                width_crop: int) -> tf.Tensor:
+    """
+    Center crop an image to the specified size.
+
+    Args:
+        image (tf.Tensor): Input image tensor.
+        output_size (Tuple): Size of image after the crop (height and width).
+
+    Returns:
+        tf.Tensor: Cropped image tensor.
+    """
+
+    # Calculate the cropping dimensions
+    input_shape = tf.shape(image)
+    height, width = input_shape[1], input_shape[2]
+
+    # Calculate the cropping offsets
+    offset_height = tf.maximum((height - height_crop) // 2, 0)
+    offset_width = tf.maximum((width - width_crop) // 2, 0)
+
+    # Crop the image
+    cropped_image = tf.image.crop_to_bounding_box(image, offset_height, offset_width, height_crop, width_crop)
+
+    return cropped_image
+
+
+def random_flip(image: tf.Tensor) -> tf.Tensor:
+    """
+    Randomly flip an image horizontally with a specified probability.
+
+    Args:
+        image (tf.Tensor): Input image tensor.
+
+    Returns:
+        tf.Tensor: Flipped image tensor.
+    """
+    flip_image = tf.image.random_flip_left_right(image)
+    return flip_image
+
+
+def clip_images(images: tf.Tensor, valid_grid: tf.Tensor, reflection: bool = False) -> tf.Tensor:
+    """
+    Clip the images based on a valid grid.
+
+    Args:
+        images (tf.Tensor): The images to be clipped.
+        valid_grid (tf.Tensor): The valid grid for clipping.
+        reflection (bool): Whether to apply reflection during clipping. Defaults to False.
+
+    Returns:
+        tf.Tensor: The clipped images.
+    """
+    clipped_images = tf.TensorArray(tf.float32, size=images.shape[1])
+
+    for i in range(valid_grid.shape[0]):
+        channel = images[:, i, :, :]
+        min_val = tf.reduce_min(valid_grid[i, :])
+        max_val = tf.reduce_max(valid_grid[i, :])
+        clamp = tf.clip_by_value(channel, min_val, max_val)
+        if reflection:
+            channel = 2 * clamp - channel
+        else:
+            channel = clamp
+        clipped_images = clipped_images.write(i, channel)
+
+    clipped_images = clipped_images.stack()
+    return tf.transpose(clipped_images, perm=[1, 0, 2, 3])
+
+
+def create_valid_grid(means, stds) -> tf.Tensor:
+    """
+    Create a valid grid for image normalization.
+
+    Returns:
+        tf.Tensor: The valid grid for image normalization.
+    """
+    # Create a pixel grid in the range 0-255, repeat for 3 color channels, and reshape
+    pixel_grid = np.arange(256).repeat(3).reshape(-1, 3)
+
+    # Transpose and add batch and channel dimensions
+    pixel_grid = tf.constant(pixel_grid, dtype=tf.float32)
+    pixel_grid = tf.transpose(pixel_grid, perm=[1, 0])
+
+    # Normalize the pixel grid using the specified mean and std
+    mean = tf.constant(np.array(means), dtype=tf.float32)
+    std = tf.constant(np.array(stds), dtype=tf.float32)
+    valid_grid = (pixel_grid - mean[: , tf.newaxis]) / std[: , tf.newaxis]
+
+    return valid_grid
+
+class Smoothing(tf.keras.layers.Layer):
+    """
+    A TensorFlow layer for applying Gaussian smoothing to an image.
+    """
+
+    def __init__(self, size: int = 3, sigma: float = 1.25):
+        """
+        Initialize the Smoothing layer.
+
+        Args:
+            size (int): The size of the Gaussian kernel.
+            sigma (float): The standard deviation of the Gaussian kernel.
+        """
+        super(Smoothing, self).__init__()
+        self.size = size
+        self.sigma = sigma
+        self.kernel = self.gaussian_kernel(size, sigma)
+
+    def build(self, input_shape):
+        """
+        Build the smoothing layer.
+
+        Args:
+            input_shape (TensorShape): Shape of the input tensor.
+        """
+        kernel = tf.reshape(self.kernel, [self.size, self.size, 1, 1])
+        self.kernel = tf.tile(kernel, [1, 1, input_shape[-1], 1])
+
+    def call(self, inputs):
+        """
+        Apply Gaussian smoothing to the input image.
+
+        Args:
+            inputs (tf.Tensor): The input image tensor.
+
+        Returns:
+            tf.Tensor: The smoothed image tensor.
+        """
+        return tf.nn.depthwise_conv2d(inputs, self.kernel, strides=[1, 1, 1, 1], padding='SAME')
+
+    def gaussian_kernel(self, size: int, sigma: float) -> tf.Tensor:
+        """
+        Create a Gaussian kernel.
+
+        Args:
+            size (int): The size of the Gaussian kernel.
+            sigma (float): The standard deviation of the Gaussian kernel.
+
+        Returns:
+            tf.Tensor: The Gaussian kernel tensor.
+        """
+        axis = tf.range(-size // 2 + 1, size // 2 + 1, dtype=tf.float32)
+        x, y = tf.meshgrid(axis, axis)
+        kernel = tf.exp(-(x ** 2 + y ** 2) / (2 * sigma ** 2))
+        kernel = kernel / tf.reduce_sum(kernel)
+        return kernel

model_compression_toolkit/data_generation/keras/image_pipeline.py

@@ -12,111 +12,62 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Tuple, Dict, Type, List
-
-import numpy as np
+from typing import Tuple, Dict, Type, Union, List
 import tensorflow as tf
 
-from model_compression_toolkit.data_generation.common.enums import ImagePipelineType, ImageNormalizationType
+from model_compression_toolkit.data_generation.common.enums import ImagePipelineType
 from model_compression_toolkit.data_generation.common.image_pipeline import BaseImagePipeline
+from model_compression_toolkit.data_generation.keras.image_operations import Smoothing, random_flip, random_crop, \
+    clip_images, create_valid_grid, center_crop
 
 
-# Define tf function for image manipulation
-
-def random_crop(image: tf.Tensor,
-                height_crop: int,
-                width_crop: int) -> tf.Tensor:
-    """
-    Randomly crop an image to the specified size.
-
-    Args:
-        image (tf.Tensor): Input image tensor.
-        height_crop (int): Size of the crop in the height axis.
-        width_crop (int): Size of the crop in the width axis.
-
-    Returns:
-        tf.Tensor: Cropped image tensor.
-    """
-    cropped_image = tf.image.random_crop(image,
-                                         size=(tf.shape(image)[0],
-                                               height_crop,
-                                               width_crop,
-                                               tf.shape(image)[-1]))
-    return cropped_image
-
-
-def center_crop(image: tf.Tensor,
-                output_size: Tuple) -> tf.Tensor:
-    """
-    Center crop an image to the specified size.
-
-    Args:
-        image (tf.Tensor): Input image tensor.
-        output_size (Tuple): Size of image after the crop (height and width).
-
-    Returns:
-        tf.Tensor: Cropped image tensor.
-    """
-
-    # Calculate the cropping dimensions
-    input_shape = tf.shape(image)
-    height, width = input_shape[1], input_shape[2]
-    target_height, target_width = output_size[0], output_size[1]
-
-    # Calculate the cropping offsets
-    offset_height = tf.maximum((height - target_height) // 2, 0)
-    offset_width = tf.maximum((width - target_width) // 2, 0)
-
-    # Crop the image
-    cropped_image = tf.image.crop_to_bounding_box(image, offset_height, offset_width, target_height, target_width)
-
-    return cropped_image
-
-
-def random_flip(image: tf.Tensor) -> tf.Tensor:
-    """
-    Randomly flip an image horizontally with a specified probability.
-
-    Args:
-        image (tf.Tensor): Input image tensor.
-
-    Returns:
-        tf.Tensor: Flipped image tensor.
-    """
-    flip_image = tf.image.random_flip_left_right(image)
-    return flip_image
-
-
-class TensorflowCropFlipImagePipeline(BaseImagePipeline):
+class TensorflowSmoothAugmentationImagePipeline(BaseImagePipeline):
     def __init__(self,
-                 output_image_size: Tuple,
-                 extra_pixels: int):
+                 output_image_size: Union[int, Tuple[int, int]],
+                 extra_pixels: Union[int, Tuple[int, int]],
+                 normalization: List[List[int]],
+                 image_clipping: bool = False,
+                 smoothing_filter_size: int = 3,
+                 smoothing_filter_sigma: float = 1.25):
         """
         Initialize the TensorflowCropFlipImagePipeline.
 
         Args:
-            output_image_size (Tuple): The output image size.
-            extra_pixels (int): Extra pixels to add to the input image size. Defaults to 0.
-        """
-        super(TensorflowCropFlipImagePipeline, self, ).__init__(output_image_size, extra_pixels)
-
+            output_image_size (Union[int, Tuple[int, int]]): The output image size.
+            extra_pixels (Union[int, Tuple[int, int]]): Extra pixels to add to the input image size. Defaults to 0.
+            normalization (List[List[float]]): The image normalization values for processing images during optimization.
+            image_clipping (bool): Whether to clip images during optimization.
+            smoothing_filter_size (int): The size of the smoothing filter. Defaults to 3.
+            smoothing_filter_sigma (float): The standard deviation of the smoothing filter. Defaults to 1.25.
+       """
+        super(TensorflowSmoothAugmentationImagePipeline, self, ).__init__(output_image_size, extra_pixels, image_clipping, normalization)
+
+        smoothing = Smoothing(smoothing_filter_size, smoothing_filter_sigma)
         # List of image manipulation functions and their arguments.
         self.img_manipulation_list = [(random_flip, {}),
-                                      (random_crop, {'height_crop': output_image_size[0],
-                                                     'width_crop': output_image_size[1]})]
+                                      (smoothing, {}),
+                                      (random_crop, {'height_crop': self.output_image_size[0],
+                                                     'width_crop': self.output_image_size[1]}),
+                                      ]
 
         # List of output image manipulation functions and their arguments.
-        self.img_output_finalize_list = [(center_crop, {'output_size': output_image_size})]
-        self.extra_pixels = extra_pixels
+        self.img_output_finalize_list = [(smoothing, {}),
+                                         (center_crop, {'height_crop': self.output_image_size[0],
+                                                        'width_crop': self.output_image_size[1]}),
+                                         ]
+        if image_clipping:
+            clip_fn = (clip_images, {'valid_grid': create_valid_grid(self.normalization[0], self.normalization[1])})
+            self.img_manipulation_list.append(clip_fn)
+            self.img_output_finalize_list.append(clip_fn)
 
-    def get_image_input_size(self) -> Tuple:
+    def get_image_input_size(self) -> Tuple[int, int]:
         """
         Get the size of the input image considering extra pixels.
 
         Returns:
-            Tuple: Size of the input image.
+            Tuple[int, int]: Size of the input image.
         """
-        return tuple(np.array(self.output_image_size) + self.extra_pixels)
+        return tuple([o + e for (o, e) in zip(self.output_image_size, self.extra_pixels)])
 
     def image_input_manipulation(self,
                                  images: tf.Tensor) -> tf.Tensor:
@@ -161,28 +112,30 @@ class TensorflowCropFlipImagePipeline(BaseImagePipeline):
 
 class TensorflowIdentityImagePipeline(BaseImagePipeline):
 
-    def __init__(self, output_image_size: int,
-                 extra_pixels: int
+    def __init__(self, output_image_size: Union[int, Tuple[int, int]],
+                 extra_pixels: Union[int, Tuple[int, int]],
+                 normalization: List[List[int]],
+                 image_clipping: bool = False
                  ):
         """
         Initialize the TensorflowIdentityImagePipeline.
 
         Args:
-            output_image_size (Tuple): The output image size.
-            extra_pixels (int): Extra pixels to add to the input image size. Defaults to 0.
+            output_image_size (Union[int, Tuple[int, int]]): The output image size.
+            extra_pixels (Union[int, Tuple[int, int]]): Extra pixels to add to the input image size. Defaults to 0.
+            normalization (List[List[float]]): The image normalization values for processing images during optimization.
+            image_clipping (bool): Whether to clip images during optimization.
         """
-        super(TensorflowIdentityImagePipeline, self, ).__init__(output_image_size, extra_pixels)
-        self.extra_pixels = extra_pixels
-        self.output_image_size = output_image_size
+        super(TensorflowIdentityImagePipeline, self, ).__init__(output_image_size, extra_pixels, image_clipping, normalization)
 
-    def get_image_input_size(self) -> Tuple:
+    def get_image_input_size(self) -> Tuple[int, int]:
         """
-        Get the size of the input image considering extra pixels.
+        Get the size of the input image.
 
         Returns:
-            Tuple: Size of the input image.
+            Tuple[int, int]: Size of the input image.
         """
-        return tuple(np.array(self.output_image_size) + self.extra_pixels)
+        return self.output_image_size
 
     def image_input_manipulation(self,
                                  images: tf.Tensor) -> tf.Tensor:
@@ -214,12 +167,5 @@ class TensorflowIdentityImagePipeline(BaseImagePipeline):
 # Dictionary mapping ImagePipelineType to corresponding image pipeline classes
 image_pipeline_dict: Dict[ImagePipelineType, Type[BaseImagePipeline]] = {
     ImagePipelineType.IDENTITY: TensorflowIdentityImagePipeline,
-    ImagePipelineType.RANDOM_CROP_FLIP: TensorflowCropFlipImagePipeline
-}
-
-# Dictionary mapping ImageNormalizationType to corresponding normalization values
-image_normalization_dict: Dict[ImageNormalizationType, List[List[float]]] = {
-    ImageNormalizationType.TORCHVISION: [[0.485, 0.456, 0.406], [0.229, 0.224, 0.225]],
-    ImageNormalizationType.KERAS_APPLICATIONS: [(127.5, 127.5, 127.5), (127.5, 127.5, 127.5)],
-    ImageNormalizationType.NO_NORMALIZATION: [[0, 0, 0], [1, 1, 1]]
+    ImagePipelineType.SMOOTHING_AND_AUGMENTATION: TensorflowSmoothAugmentationImagePipeline
 }

model_compression_toolkit/data_generation/keras/keras_data_generation.py

@@ -13,12 +13,13 @@
 # limitations under the License.
 # ==============================================================================
 import time
-from typing import Callable, Tuple, List, Dict
+from typing import Callable, Tuple, List, Dict, Union
 from tqdm import tqdm
 
 from model_compression_toolkit.constants import FOUND_TF
 from model_compression_toolkit.data_generation.common.constants import DEFAULT_N_ITER, DEFAULT_DATA_GEN_BS
 from model_compression_toolkit.data_generation.common.data_generation import get_data_generation_classes
+from model_compression_toolkit.data_generation.common.image_pipeline import image_normalization_dict
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.data_generation.common.data_generation_config import DataGenerationConfig, \
     ImageGranularity
@@ -30,9 +31,8 @@ if FOUND_TF:
     from tensorflow.keras.layers import BatchNormalization
     from tensorflow.keras.optimizers.legacy import Optimizer, Adam
     from model_compression_toolkit.data_generation.keras.constants import DEFAULT_KERAS_INITIAL_LR, \
-        DEFAULT_KERAS_OUTPUT_LOSS_MULTIPLIER
-    from model_compression_toolkit.data_generation.keras.image_pipeline import (image_pipeline_dict,
-                                                                                image_normalization_dict)
+    DEFAULT_KERAS_EXTRA_PIXELS, DEFAULT_KERAS_OUTPUT_LOSS_MULTIPLIER
+    from model_compression_toolkit.data_generation.keras.image_pipeline import image_pipeline_dict
     from model_compression_toolkit.data_generation.keras.model_info_exctractors import (KerasActivationExtractor,
                                                                                         KerasOriginalBNStatsHolder)
     from model_compression_toolkit.data_generation.keras.optimization_functions.batchnorm_alignment_functions import \
@@ -55,18 +55,17 @@ if FOUND_TF:
             data_gen_batch_size: int = DEFAULT_DATA_GEN_BS,
             initial_lr: float = DEFAULT_KERAS_INITIAL_LR,
             output_loss_multiplier: float = DEFAULT_KERAS_OUTPUT_LOSS_MULTIPLIER,
-            scheduler_type: SchedulerType = SchedulerType.REDUCE_ON_PLATEAU,
+            scheduler_type: SchedulerType = SchedulerType.REDUCE_ON_PLATEAU   ,
             bn_alignment_loss_type: BatchNormAlignemntLossType = BatchNormAlignemntLossType.L2_SQUARE,
             output_loss_type: OutputLossType = OutputLossType.REGULARIZED_MIN_MAX_DIFF,
             data_init_type: DataInitType = DataInitType.Gaussian,
             layer_weighting_type: BNLayerWeightingType = BNLayerWeightingType.AVERAGE,
             image_granularity: ImageGranularity = ImageGranularity.BatchWise,
-            image_pipeline_type: ImagePipelineType = ImagePipelineType.RANDOM_CROP_FLIP,
+            image_pipeline_type: ImagePipelineType = ImagePipelineType.SMOOTHING_AND_AUGMENTATION,
             image_normalization_type: ImageNormalizationType = ImageNormalizationType.KERAS_APPLICATIONS,
-            extra_pixels: int = 0,
+            extra_pixels: Union[int, Tuple[int, int]] = DEFAULT_KERAS_EXTRA_PIXELS,
             bn_layer_types: List = [BatchNormalization],
-            clip_images: bool = True,
-            reflection: bool = True,
+            image_clipping: bool = False,
     ) -> DataGenerationConfig:
         """
         Function to create a DataGenerationConfig object with the specified configuration parameters.
@@ -85,10 +84,9 @@ if FOUND_TF:
             image_granularity (ImageGranularity): The granularity of the images for optimization.
             image_pipeline_type (ImagePipelineType): The type of image pipeline to use.
             image_normalization_type (ImageNormalizationType): The type of image normalization to use.
-            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 to be considered for data generation.
-            clip_images (bool): Whether to clip images during optimization.
-            reflection (bool): Whether to use reflection during optimization.
+            image_clipping (bool): Whether to clip images during optimization.
 
         Returns:
             DataGenerationConfig: Data generation configuration object.
@@ -100,6 +98,7 @@ if FOUND_TF:
             optimizer=optimizer,
             data_gen_batch_size=data_gen_batch_size,
             initial_lr=initial_lr,
+            output_loss_multiplier=output_loss_multiplier,
             scheduler_type=scheduler_type,
             bn_alignment_loss_type=bn_alignment_loss_type,
             output_loss_type=output_loss_type,
@@ -110,15 +109,13 @@ if FOUND_TF:
             image_normalization_type=image_normalization_type,
             extra_pixels=extra_pixels,
             bn_layer_types=bn_layer_types,
-            clip_images=clip_images,
-            reflection=reflection,
-            output_loss_multiplier=output_loss_multiplier)
+            image_clipping=image_clipping)
 
 
     def keras_data_generation_experimental(
             model: tf.keras.Model,
             n_images: int,
-            output_image_size: Tuple,
+            output_image_size: Union[int, Tuple[int, int]],
             data_generation_config: DataGenerationConfig) -> tf.Tensor:
         """
         Function to perform data generation using the provided Keras model and data generation configuration.
@@ -126,7 +123,7 @@ if FOUND_TF:
         Args:
             model (Model): Keras model to generate data for.
             n_images (int): Number of images to generate.
-            output_image_size (Tuple): Size of the output images.
+            output_image_size (Union[int, Tuple[int, int]]): Size of the output images.
             data_generation_config (DataGenerationConfig): Configuration for data generation.
 
         Returns:
@@ -180,17 +177,13 @@ if FOUND_TF:
                                                        bn_alignment_loss_function_dict=bn_alignment_loss_function_dict,
                                                        output_loss_function_dict=output_loss_function_dict)
 
-        if not all(normalization[1]):
-            Logger.critical(
-                f'Invalid normalization standard deviation {normalization[1]} set to zero, which will lead to division by zero. Please select a non-zero normalization standard deviation.')
-
         # Get the scheduler functions corresponding to the specified scheduler type
         scheduler_get_fn = scheduler_step_function_dict.get(data_generation_config.scheduler_type)
 
         # Check if the scheduler type is valid
         if scheduler_get_fn is None:
             Logger.critical(
-                f'Invalid scheduler_type {data_generation_config.scheduler_type}. Please select one from {SchedulerType.get_values()}.')
+                f'Invalid scheduler_type {data_generation_config.scheduler_type}. Please select one from {SchedulerType.get_values()}.') # pragma: no cover
 
         # Create a scheduler object with the specified number of iterations
         scheduler = scheduler_get_fn(n_iter=data_generation_config.n_iter,
@@ -202,10 +195,7 @@ if FOUND_TF:
         # Create an activation extractor object to extract activations from the model
         activation_extractor = KerasActivationExtractor(model=model,
                                                         layer_types_to_extract_inputs=
-                                                        data_generation_config.bn_layer_types,
-                                                        image_granularity=data_generation_config.image_granularity,
-                                                        image_input_manipulation=
-                                                        image_pipeline.image_input_manipulation)
+                                                        data_generation_config.bn_layer_types)
 
         # Create an orig_bn_stats_holder object to hold original BatchNorm statistics
         orig_bn_stats_holder = KerasOriginalBNStatsHolder(model=model,
@@ -223,9 +213,6 @@ if FOUND_TF:
             model=model,
             orig_bn_stats_holder=orig_bn_stats_holder)
 
-        # Compute the layer weights based on orig_bn_stats_holder
-        bn_layer_weights = bn_layer_weighting_fn(orig_bn_stats_holder=orig_bn_stats_holder)
-
         # Get the current time to measure the total time taken
         total_time = time.time()
 
@@ -233,7 +220,7 @@ if FOUND_TF:
         ibar = tqdm(range(data_generation_config.n_iter))
 
         # Perform data generation iterations
-        for i_ter in ibar:
+        for i_iter in ibar:
 
             # Randomly reorder the batches
             all_imgs_opt_handler.random_batch_reorder()
@@ -246,6 +233,12 @@ if FOUND_TF:
                 # Get the images to optimize and the optimizer for the batch
                 imgs_to_optimize = all_imgs_opt_handler.get_images_by_batch_index(batch_index=random_batch_index)
 
+                # Compute the layer weights based on orig_bn_stats_holder
+                bn_layer_weights = bn_layer_weighting_fn(orig_bn_stats_holder=orig_bn_stats_holder,
+                                                         activation_extractor=activation_extractor,
+                                                         i_iter=i_iter,
+                                                         n_iter=data_generation_config.n_iter)
+
                 # Compute the gradients and the loss for the batch
                 gradients, total_loss, bn_loss, output_loss = keras_compute_grads(imgs_to_optimize=imgs_to_optimize,
                                                                                   batch_index=random_batch_index,
@@ -266,7 +259,7 @@ if FOUND_TF:
                                                        images=imgs_to_optimize,
                                                        gradients=gradients,
                                                        loss=total_loss,
-                                                       i_ter=i_ter)
+                                                       i_iter=i_iter)
 
                 # Update the statistics based on the updated images
                 if all_imgs_opt_handler.use_all_data_stats:
@@ -335,14 +328,13 @@ if FOUND_TF:
                                                            bn_layer_weights=bn_layer_weights)
 
             # Compute output loss
-            # If output_loss_multiplier is zero return 0
-            output_loss = output_loss_multiplier * output_loss_fn(
-                output_imgs=output,
+            output_loss = output_loss_fn(
+                model_outputs=output,
                 activation_extractor=activation_extractor,
-                tape=tape) if output_loss_multiplier > 0 else tf.zeros(1)
+                tape=tape)
 
             # Compute total loss
-            total_loss = bn_loss + output_loss
+            total_loss = bn_loss + output_loss_multiplier * output_loss
 
             # Get the trainable variables
             variables = [imgs_to_optimize]

model_compression_toolkit/data_generation/keras/model_info_exctractors.py

@@ -72,8 +72,6 @@ class KerasActivationExtractor(ActivationExtractor):
     def __init__(self,
                  model: tf.keras.Model,
                  layer_types_to_extract_inputs: List,
-                 image_granularity: ImageGranularity,
-                 image_input_manipulation: Callable,
                  linear_layers: Tuple = (Dense, Conv2D)):
         """
         Initializes the KerasActivationExtractor.
@@ -81,14 +79,10 @@ class KerasActivationExtractor(ActivationExtractor):
         Args:
             model (Model): Keras model to generate data for.
             layer_types_to_extract_inputs (List): Tuple or list of layer types.
-            image_granularity (ImageGranularity): The granularity of the images for optimization.
-            image_input_manipulation (Callable): Function for image input manipulation.
             linear_layers (Tuple): Tuple of linear layers types to retrieve the output of the last linear layer
 
         """
         self.model = model
-        self.image_input_manipulation = image_input_manipulation
-        self.image_granularity = image_granularity
         self.layer_types_to_extract_inputs = tuple(layer_types_to_extract_inputs)
         self.linear_layers = linear_layers
 
@@ -96,7 +90,6 @@ class KerasActivationExtractor(ActivationExtractor):
         self.bn_layer_names = [layer.name for layer in model.layers if isinstance(layer,
                                                                                   self.layer_types_to_extract_inputs)]
         self.num_layers = len(self.bn_layer_names)
-        Logger.info(f'Number of layers = {self.num_layers}')
 
         # Initialize stats containers
         self.activations = {}
@@ -206,9 +199,3 @@ class KerasActivationExtractor(ActivationExtractor):
                     last_layer = layer
                     break
         return last_layer
-
-    def remove(self):
-        """
-        Remove the stats containers.
-        """
-        self.activations = {}

model_compression_toolkit/data_generation/keras/optimization_functions/bn_layer_weighting_functions.py

@@ -15,16 +15,22 @@
 from typing import Dict, Callable
 
 from model_compression_toolkit.data_generation.common.enums import BNLayerWeightingType
-from model_compression_toolkit.data_generation.keras.model_info_exctractors import KerasOriginalBNStatsHolder
+from model_compression_toolkit.data_generation.keras.model_info_exctractors import KerasOriginalBNStatsHolder, \
+    KerasActivationExtractor
 
 
-def average_layer_weighting_fn(orig_bn_stats_holder: KerasOriginalBNStatsHolder, **kwargs) -> Dict[str, float]:
+def average_layer_weighting_fn(orig_bn_stats_holder: KerasOriginalBNStatsHolder,
+                               activation_extractor: KerasActivationExtractor,
+                               i_iter: int,
+                               n_iter: int) -> Dict[str, float]:
     """
     Calculate average weighting for each batch normalization layer.
 
     Args:
         orig_bn_stats_holder (KerasOriginalBNStatsHolder): Holder for original batch normalization statistics.
-        **kwargs: Additional arguments if needed.
+        activation_extractor (KerasActivationExtractor): The activation extractor for the model.
+        i_iter (int): Current optimization iteration.
+        n_iter (int): Total number of optimization iterations.
 
     Returns:
         Dict[str, float]: A dictionary containing layer names as keys and average weightings as values.
@@ -33,14 +39,18 @@ def average_layer_weighting_fn(orig_bn_stats_holder: KerasOriginalBNStatsHolder,
     return {bn_layer_name: 1 / num_bn_layers for bn_layer_name in orig_bn_stats_holder.get_bn_layer_names()}
 
 
-def first_bn_multiplier_weighting_fn(orig_bn_stats_holder: KerasOriginalBNStatsHolder, **kwargs) -> Dict[str, float]:
+def first_bn_multiplier_weighting_fn(orig_bn_stats_holder: KerasOriginalBNStatsHolder,
+                                     activation_extractor: KerasActivationExtractor,
+                                     i_iter: int,
+                                     n_iter: int) -> Dict[str, float]:
     """
     Calculate layer weightings with a higher multiplier for the first batch normalization layer.
 
     Args:
         orig_bn_stats_holder (KerasOriginalBNStatsHolder): Holder for original batch normalization statistics.
-        **kwargs: Additional arguments if needed.
-
+        activation_extractor (KerasActivationExtractor): The activation extractor for the model.
+        i_iter (int): Current optimization iteration.
+        n_iter (int): Total number of optimization iterations.
     Returns:
         Dict[str, float]: A dictionary containing layer names as keys and weightings as values.
     """