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/pytorch/optimization_functions/output_loss_functions.py

@@ -17,37 +17,66 @@ from typing import Dict, Callable
 import torch
 from torch import Tensor
 
-from model_compression_toolkit.core.pytorch.pytorch_device_config import get_working_device
 from model_compression_toolkit.data_generation.common.enums import OutputLossType
 from model_compression_toolkit.data_generation.pytorch.model_info_exctractors import PytorchActivationExtractor
 
-def min_max_diff(
-        output_imgs: Tensor,
+def inverse_min_max_diff(
+        model_outputs: Tensor,
         activation_extractor: PytorchActivationExtractor,
+        device: torch.device,
         eps: float = 1e-6) -> Tensor:
     """
-    Calculate the minimum-maximum difference of output images.
+    Calculate the inverse of the maximum - minimum difference of the model output on the input images.
 
     Args:
-        output_imgs (Tensor or List[Tensor]): The output of the model on images.
+        model_outputs (Tensor or List[Tensor]): The output of the model on images.
         activation_extractor (PytorchActivationExtractor): The activation extractor for the model.
+        device (torch.device): The current device set for PyTorch operations.
         eps (float): Small value for numerical stability.
 
     Returns:
         Tensor: The computed minimum-maximum difference loss.
     """
-    if not isinstance(output_imgs, (list, tuple)):
-        output_imgs = [output_imgs]
-    output_loss = 0
-    for output in output_imgs:
+    if not isinstance(model_outputs, (list, tuple)):
+        model_outputs = [model_outputs]
+    output_loss = torch.zeros(1).to(device)
+    for output in model_outputs:
         output = torch.reshape(output, [output.shape[0], -1])
         output_loss += 1 / torch.mean(torch.max(output, 1)[0] - torch.min(output, 1)[0] + eps)
     return output_loss
 
+def negative_min_max_diff(
+        model_outputs: Tensor,
+        activation_extractor: PytorchActivationExtractor,
+        device: torch.device,
+        eps: float = 1e-6) -> Tensor:
+    """
+    Calculate the mean of the negative maximum - minimum difference of the model output on the input images.
+
+    Args:
+        model_outputs (Tensor or List[Tensor]): The output of the model on images.
+        activation_extractor (PytorchActivationExtractor): The activation extractor for the model.
+        device (torch.device): The current device set for PyTorch operations.
+        eps (float): Small value for numerical stability.
+
+    Returns:
+        Tensor: The computed minimum-maximum difference loss.
+    """
+    if not isinstance(model_outputs, (list, tuple)):
+        model_outputs = [model_outputs]
+    output_loss = torch.zeros(1).to(device)
+    for output in model_outputs:
+        output = torch.reshape(output, [output.shape[0], -1])
+        out_max, out_argmax = torch.max(output, dim=1)
+        out_min, out_argmin = torch.min(output, dim=1)
+        output_loss += torch.mean(-(out_max - out_min))
+    return output_loss
+
 
 def regularized_min_max_diff(
-        output_imgs: Tensor,
+        model_outputs: Tensor,
         activation_extractor: PytorchActivationExtractor,
+        device: torch.device,
         eps: float = 1e-6) -> Tensor:
     """
     Calculate the regularized minimum-maximum difference of output images. We want to maximize
@@ -56,8 +85,9 @@ def regularized_min_max_diff(
     the last layer's weights.
 
     Args:
-        output_imgs (Tensor or List[Tensor]): The output of the model on images.
+        model_outputs (Tensor or List[Tensor]): The output of the model on images.
         activation_extractor (PytorchActivationExtractor): The activation extractor for the model.
+        device (torch.device): The current device set for PyTorch operations.
         eps (float): Small value for numerical stability.
 
     Returns:
@@ -69,13 +99,13 @@ def regularized_min_max_diff(
     # get the weights of the last linear layers of the model
     weights_output_layers = activation_extractor.get_last_linear_layers_weights()
 
-    if not isinstance(output_imgs, (list, tuple)):
-        output_imgs = torch.reshape(output_imgs, [output_imgs.shape[0], output_imgs.shape[1], -1])
-        output_imgs = torch.mean(output_imgs, dim=-1)
-        output_imgs = [output_imgs]
-    output_loss = torch.zeros(1).to(get_working_device())
+    if not isinstance(model_outputs, (list, tuple)):
+        model_outputs = torch.reshape(model_outputs, [model_outputs.shape[0], model_outputs.shape[1], -1])
+        model_outputs = torch.mean(model_outputs, dim=-1)
+        model_outputs = [model_outputs]
+    output_loss = torch.zeros(1).to(device)
 
-    for output_weight, output, last_layer_input in zip(weights_output_layers, output_imgs, output_layers_inputs):
+    for output_weight, output, last_layer_input in zip(weights_output_layers, model_outputs, output_layers_inputs):
         weights_norm = torch.linalg.norm(output_weight.squeeze(), dim=1)
         out_max, out_argmax = torch.max(output, dim=1)
         out_min, out_argmin = torch.min(output, dim=1)
@@ -88,27 +118,31 @@ def regularized_min_max_diff(
         output_loss += torch.mean(reg_min + reg_max + dynamic_loss)
     return output_loss
 
+
 def no_output_loss(
-        output_imgs: Tensor,
+        model_outputs: Tensor,
         activation_extractor: PytorchActivationExtractor,
+        device: torch.device,
         eps: float = 1e-6) -> Tensor:
     """
     Calculate no output loss.
 
     Args:
-        output_imgs (Tensor): The output of the model on images.
+        model_outputs (Tensor): The output of the model on images.
         activation_extractor (PytorchActivationExtractor): The activation extractor for the model.
+        device (torch.device): The current device set for PyTorch operations.
         eps (float): Small value for numerical stability.
 
     Returns:
         Tensor: A tensor with zero value for the loss.
     """
-    return torch.zeros(1).to(get_working_device())
+    return torch.zeros(1).to(device)
 
 
 # Dictionary of output loss functions
 output_loss_function_dict: Dict[OutputLossType, Callable] = {
     OutputLossType.NONE: no_output_loss,
-    OutputLossType.MIN_MAX_DIFF: min_max_diff,
+    OutputLossType.NEGATIVE_MIN_MAX_DIFF: negative_min_max_diff,
+    OutputLossType.INVERSE_MIN_MAX_DIFF: inverse_min_max_diff,
     OutputLossType.REGULARIZED_MIN_MAX_DIFF: regularized_min_max_diff,
 }

model_compression_toolkit/data_generation/pytorch/optimization_functions/scheduler_step_functions.py

@@ -17,6 +17,8 @@ from typing import Callable, Any, Dict, Tuple
 
 from torch.optim.lr_scheduler import ReduceLROnPlateau, StepLR
 from model_compression_toolkit.data_generation.common.enums import SchedulerType
+from model_compression_toolkit.data_generation.pytorch.optimization_functions.lr_scheduler import \
+    ReduceLROnPlateauWithReset
 
 
 def get_reduce_lr_on_plateau_scheduler(n_iter: int) -> Callable:
@@ -31,6 +33,18 @@ def get_reduce_lr_on_plateau_scheduler(n_iter: int) -> Callable:
     """
     return partial(ReduceLROnPlateau, min_lr=1e-4, factor=0.5, patience=int(n_iter / 50))
 
+def get_reduce_lr_on_plateau_with_reset_scheduler(n_iter: int) -> Callable:
+    """
+    Get a ReduceLROnPlateauWithReset scheduler.
+
+    Args:
+        n_iter (int): The number of iterations.
+
+    Returns:
+        Callable: A partial function to create ReduceLROnPlateauWithReset scheduler with specified parameters.
+    """
+    return partial(ReduceLROnPlateauWithReset, min_lr=1e-4, factor=0.5, patience=int(n_iter / 50))
+
 def get_step_lr_scheduler(n_iter: int) -> Callable:
     """
     Get a StepLR scheduler.
@@ -69,5 +83,6 @@ def scheduler_step_fn(scheduler: Any, i_iter: int, loss_value: float):
 # Dictionary of scheduler functions and their corresponding step functions
 scheduler_step_function_dict: Dict[SchedulerType, Tuple[Callable, Callable]] = {
     SchedulerType.REDUCE_ON_PLATEAU: (get_reduce_lr_on_plateau_scheduler, reduce_lr_on_platu_step_fn),
+    SchedulerType.REDUCE_ON_PLATEAU_WITH_RESET: (get_reduce_lr_on_plateau_with_reset_scheduler, reduce_lr_on_platu_step_fn),
     SchedulerType.STEP: (get_step_lr_scheduler, scheduler_step_fn),
 }

model_compression_toolkit/data_generation/pytorch/optimization_utils.py

@@ -20,15 +20,16 @@ from torch import Tensor
 from torch.nn import Module
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader, Dataset
-from torchvision.transforms import Normalize
+from torch.cuda.amp import GradScaler
 
-from model_compression_toolkit.core.pytorch.pytorch_device_config import get_working_device
+from model_compression_toolkit.core.pytorch.utils import to_torch_tensor, clip_inf_values_float16
 from model_compression_toolkit.data_generation.common.enums import ImageGranularity
 from model_compression_toolkit.data_generation.common.image_pipeline import BaseImagePipeline
 from model_compression_toolkit.data_generation.common.optimization_utils import BatchStatsHolder, AllImagesStatsHolder, \
     BatchOptimizationHolder, ImagesOptimizationHandler
 from model_compression_toolkit.data_generation.common.constants import IMAGE_INPUT
 from model_compression_toolkit.data_generation.pytorch.constants import BATCH_AXIS, H_AXIS, W_AXIS
+from model_compression_toolkit.data_generation.pytorch.image_operations import create_valid_grid
 from model_compression_toolkit.data_generation.pytorch.model_info_exctractors import ActivationExtractor
 
 
@@ -58,8 +59,7 @@ class PytorchImagesOptimizationHandler(ImagesOptimizationHandler):
                  initial_lr: float,
                  normalization_mean: List[float],
                  normalization_std: List[float],
-                 clip_images: bool,
-                 reflection: bool,
+                 device: str,
                  eps: float = 1e-6):
         """
         Constructor for the PytorchImagesOptimizationHandler class.
@@ -77,8 +77,7 @@ class PytorchImagesOptimizationHandler(ImagesOptimizationHandler):
             initial_lr (float): The initial learning rate used by the optimizer.
             normalization_mean (List[float]): The mean values for image normalization.
             normalization_std (List[float]): The standard deviation values for image normalization.
-            clip_images (bool): Whether to clip the images during optimization.
-            reflection (bool): Whether to use reflection during image clipping.
+            device (torch.device): The current device set for PyTorch operations.
             eps (float): A small value added for numerical stability.
         """
         super(PytorchImagesOptimizationHandler, self).__init__(model=model,
@@ -93,16 +92,12 @@ class PytorchImagesOptimizationHandler(ImagesOptimizationHandler):
                                                                   initial_lr=initial_lr,
                                                                   normalization_mean=normalization_mean,
                                                                   normalization_std=normalization_std,
-                                                                  clip_images=clip_images,
-                                                                  reflection=reflection,
                                                                   eps=eps)
 
-        self.device = get_working_device()
-        # Image valid grid, each image value can only be 0 - 255 before normalization
-        t = torch.from_numpy(np.array(list(range(256))).repeat(3).reshape(-1, 3) / 255)
-        self.valid_grid = Normalize(mean=normalization_mean,
-                                    std=normalization_std)(t.transpose(1, 0)[None, :, :, None]).squeeze().to(self.device)
-
+        # Initialize mixed-precision scaler
+        self.scaler = GradScaler()
+        self.device = device
+        self.valid_grid = create_valid_grid(normalization_mean, normalization_std)
 
         # Set the mean axis based on the image granularity
         if self.image_granularity == ImageGranularity.ImageWise:
@@ -155,42 +150,40 @@ class PytorchImagesOptimizationHandler(ImagesOptimizationHandler):
         total_mean, total_second_moment = 0, 0
         for i_batch in range(self.n_batches):
             mean, second_moment, std = self.all_imgs_stats_holder.get_stats(i_batch, layer_name)
-            total_mean += mean
-            total_second_moment += second_moment
+            if mean is not None:
+                total_mean += mean
+            if second_moment is not None:
+                total_second_moment += second_moment
 
         total_mean /= self.n_batches
         total_second_moment /= self.n_batches
-        total_var = total_second_moment - torch.pow(total_mean, 2)
+        total_var = to_torch_tensor(total_second_moment) - torch.pow(to_torch_tensor(total_mean), 2)
         total_std = torch.sqrt(total_var + self.eps)
         return total_mean, total_std
 
     def optimization_step(self,
                           batch_index: int,
                           loss: Tensor,
-                          i_ter: int):
+                          i_iter: int):
         """
         Perform an optimization step.
 
         Args:
             batch_index (int): Index of the batch.
             loss (Tensor): Loss value.
-            i_ter (int): Current optimization iteration.
+            i_iter (int): Current optimization iteration.
         """
         # Get optimizer and scheduler for the specific batch index
         optimizer = self.get_optimizer_by_batch_index(batch_index)
         scheduler = self.get_scheduler_by_batch_index(batch_index)
 
         # Backward pass
-        loss.backward()
-
-        # Update weights
-        optimizer.step()
+        self.scaler.scale(loss).backward()
+        self.scaler.step(optimizer)
+        self.scaler.update()
 
         # Perform scheduler step
-        self.scheduler_step_fn(scheduler, i_ter, loss.item())
-
-        if self.clip_images:
-            self.batch_opt_holders_list[batch_index].clip_images(self.valid_grid, reflection=self.reflection)
+        self.scheduler_step_fn(scheduler, i_iter, loss.item())
 
 
     def zero_grad(self, batch_index: int):
@@ -259,25 +252,6 @@ class PytorchBatchOptimizationHolder(BatchOptimizationHolder):
         self.optimizer = optimizer([self.images], lr=initial_lr)
         self.scheduler = scheduler(self.optimizer)
 
-    def clip_images(self,
-                    valid_grid: Tensor,
-                    reflection: bool = True):
-        """
-        Clip the images.
-
-        Args:
-            valid_grid (Tensor): A tensor containing valid values for image clipping.
-            reflection (bool): Whether to use reflection during image clipping. Defaults to True.
-        """
-        with torch.no_grad():
-            for i_ch in range(valid_grid.shape[0]):
-                clamp = torch.clamp(self.images[:, i_ch, :, :], valid_grid[i_ch, :].min(), valid_grid[i_ch, :].max())
-                if reflection:
-                    self.images[:, i_ch, :, :] = 2 * clamp - self.images[:, i_ch, :, :]
-                else:
-                    self.images[:, i_ch, :, :] = clamp
-        self.images.requires_grad = True
-
 
 class PytorchAllImagesStatsHolder(AllImagesStatsHolder):
     """
@@ -332,8 +306,9 @@ class PytorchBatchStatsHolder(BatchStatsHolder):
         """
         mean = self.get_mean(bn_layer_name)
         second_moment = self.get_second_moment(bn_layer_name)
-        var = second_moment - torch.pow(mean, 2.0)
-        return var
+        if mean is not None and second_moment is not None:
+            return second_moment - torch.pow(mean, 2.0)
+        return None
 
 
     def get_std(self, bn_layer_name: str) -> Tensor:
@@ -347,7 +322,9 @@ class PytorchBatchStatsHolder(BatchStatsHolder):
             Tensor: The standard deviation for the specified layer.
         """
         var = self.get_var(bn_layer_name)
-        return torch.sqrt(var + self.eps)
+        if var is not None:
+            return torch.sqrt(var + self.eps)
+        return None
 
     def calc_bn_stats_from_activations(self,
                                        input_imgs: Tensor,
@@ -374,12 +351,13 @@ class PytorchBatchStatsHolder(BatchStatsHolder):
         # Extract statistics of intermediate convolution outputs before the BatchNorm layers
         for bn_layer_name in activation_extractor.get_extractor_layer_names():
             bn_input_activations = activation_extractor.get_layer_input_activation(bn_layer_name)
-            if not to_differentiate:
-                bn_input_activations = bn_input_activations.detach()
+            if bn_input_activations is not None:
+                if not to_differentiate:
+                    bn_input_activations = bn_input_activations.detach()
 
-            collected_mean = torch.mean(bn_input_activations, dim=self.mean_axis)
-            collected_second_moment = torch.mean(torch.pow(bn_input_activations, 2.0), dim=self.mean_axis)
-            self.update_layer_stats(bn_layer_name, collected_mean, collected_second_moment)
+                collected_mean = torch.mean(bn_input_activations, dim=self.mean_axis)
+                collected_second_moment = clip_inf_values_float16(torch.mean(torch.pow(bn_input_activations, 2.0), dim=self.mean_axis))
+                self.update_layer_stats(bn_layer_name, collected_mean, collected_second_moment)
 
     def clear(self):
         """Clear the statistics."""

model_compression_toolkit/data_generation/pytorch/pytorch_data_generation.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 import time
-from typing import Callable, Any, Tuple, List
+from typing import Callable, Any, Tuple, List, Union
 
 from tqdm import tqdm
 
@@ -25,10 +25,12 @@ from model_compression_toolkit.data_generation.common.data_generation_config imp
 from model_compression_toolkit.data_generation.common.enums import ImageGranularity, SchedulerType, \
     BatchNormAlignemntLossType, DataInitType, BNLayerWeightingType, ImagePipelineType, ImageNormalizationType, \
     OutputLossType
+from model_compression_toolkit.data_generation.common.image_pipeline import image_normalization_dict
 from model_compression_toolkit.data_generation.pytorch.constants import DEFAULT_PYTORCH_INITIAL_LR, \
-    DEFAULT_PYTORCH_OUTPUT_LOSS_MULTIPLIER, DEFAULT_PYTORCH_BN_LAYER_TYPES, DEFAULT_PYTORCH_LAST_LAYER_TYPES
+    DEFAULT_PYTORCH_BN_LAYER_TYPES, DEFAULT_PYTORCH_LAST_LAYER_TYPES, DEFAULT_PYTORCH_EXTRA_PIXELS, \
+    DEFAULT_PYTORCH_OUTPUT_LOSS_MULTIPLIER
 from model_compression_toolkit.data_generation.pytorch.image_pipeline import image_pipeline_dict, \
-    image_normalization_dict, BaseImagePipeline
+    BaseImagePipeline
 from model_compression_toolkit.data_generation.pytorch.model_info_exctractors import PytorchActivationExtractor, \
     PytorchOriginalBNStatsHolder
 from model_compression_toolkit.data_generation.pytorch.optimization_functions.batchnorm_alignment_functions import \
@@ -51,6 +53,7 @@ if FOUND_TORCH and FOUND_TORCHVISION:
     from torch.nn import Module
     from torch.optim import RAdam, Optimizer
     from torch.fx import symbolic_trace
+    from torch.cuda.amp import autocast
 
     from model_compression_toolkit.core.pytorch.pytorch_device_config import get_working_device
 
@@ -61,19 +64,18 @@ if FOUND_TORCH and FOUND_TORCHVISION:
             data_gen_batch_size=DEFAULT_DATA_GEN_BS,
             initial_lr=DEFAULT_PYTORCH_INITIAL_LR,
             output_loss_multiplier=DEFAULT_PYTORCH_OUTPUT_LOSS_MULTIPLIER,
-            scheduler_type: SchedulerType = SchedulerType.REDUCE_ON_PLATEAU,
+            scheduler_type: SchedulerType = SchedulerType.REDUCE_ON_PLATEAU_WITH_RESET,
             bn_alignment_loss_type: BatchNormAlignemntLossType = BatchNormAlignemntLossType.L2_SQUARE,
-            output_loss_type: OutputLossType = OutputLossType.REGULARIZED_MIN_MAX_DIFF,
-            data_init_type: DataInitType = DataInitType.Diverse,
+            output_loss_type: OutputLossType = OutputLossType.NEGATIVE_MIN_MAX_DIFF,
+            data_init_type: DataInitType = DataInitType.Gaussian,
             layer_weighting_type: BNLayerWeightingType = BNLayerWeightingType.AVERAGE,
             image_granularity=ImageGranularity.AllImages,
-            image_pipeline_type: ImagePipelineType = ImagePipelineType.RANDOM_CROP,
+            image_pipeline_type: ImagePipelineType = ImagePipelineType.SMOOTHING_AND_AUGMENTATION,
             image_normalization_type: ImageNormalizationType = ImageNormalizationType.TORCHVISION,
-            extra_pixels: int = 0,
+            extra_pixels: Union[int, Tuple[int, int]] = DEFAULT_PYTORCH_EXTRA_PIXELS,
             bn_layer_types: List = DEFAULT_PYTORCH_BN_LAYER_TYPES,
             last_layer_types: List = DEFAULT_PYTORCH_LAST_LAYER_TYPES,
-            clip_images: bool = True,
-            reflection: bool = True,
+            image_clipping: bool = True,
     ) -> DataGenerationConfig:
         """
         Function to create a DataGenerationConfig object with the specified configuration parameters.
@@ -92,11 +94,10 @@ if FOUND_TORCH and FOUND_TORCHVISION:
             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.
             last_layer_types (List): List of layer types to be considered for the output loss.
-            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:
@@ -121,15 +122,14 @@ if FOUND_TORCH and FOUND_TORCHVISION:
             extra_pixels=extra_pixels,
             bn_layer_types=bn_layer_types,
             last_layer_types=last_layer_types,
-            clip_images=clip_images,
-            reflection=reflection
+            image_clipping=image_clipping,
         )
 
 
     def pytorch_data_generation_experimental(
             model: Module,
             n_images: int,
-            output_image_size: int,
+            output_image_size: Union[int, Tuple[int, int]],
             data_generation_config: DataGenerationConfig) -> List[Tensor]:
         """
         Function to perform data generation using the provided model and data generation configuration.
@@ -137,7 +137,7 @@ if FOUND_TORCH and FOUND_TORCHVISION:
         Args:
             model (Module): PyTorch model to generate data for.
             n_images (int): Number of images to generate.
-            output_image_size (int): The hight and width size of the output images.
+            output_image_size (Union[int, Tuple[int, int]]): The hight and width size of the output images.
             data_generation_config (DataGenerationConfig): Configuration for data generation.
 
         Returns:
@@ -176,6 +176,9 @@ if FOUND_TORCH and FOUND_TORCHVISION:
                        f"If you encounter an issue, please open an issue in our GitHub "
                        f"project https://github.com/sony/model_optimization")
 
+        # get the model device
+        device = get_working_device()
+
         # get a static graph representation of the model using torch.fx
         fx_model = symbolic_trace(model)
 		
@@ -198,8 +201,8 @@ if FOUND_TORCH and FOUND_TORCHVISION:
 
         # Check if the scheduler type is valid
         if scheduler_get_fn is None or scheduler_step_fn is None:
-            Logger.critical(f'Invalid output_loss_type {data_generation_config.scheduler_type}. '
-                            f'Please select one from {SchedulerType.get_values()}.')
+            Logger.critical(f'Invalid scheduler_type {data_generation_config.scheduler_type}. '
+                            f'Please select one from {SchedulerType.get_values()}.') # pragma: no cover
 
         # Create a scheduler object with the specified number of iterations
         scheduler = scheduler_get_fn(data_generation_config.n_iter)
@@ -218,23 +221,22 @@ if FOUND_TORCH and FOUND_TORCHVISION:
         orig_bn_stats_holder = PytorchOriginalBNStatsHolder(model, data_generation_config.bn_layer_types)
         if orig_bn_stats_holder.get_num_bn_layers() == 0:
             Logger.critical(
-                f'Data generation requires a model with at least one BatchNorm layer.')
+                f'Data generation requires a model with at least one BatchNorm layer.') # pragma: no cover
 
         # Create an ImagesOptimizationHandler object for handling optimization
         all_imgs_opt_handler = PytorchImagesOptimizationHandler(model=model,
-                                                                   data_gen_batch_size=data_generation_config.data_gen_batch_size,
-                                                                   init_dataset=init_dataset,
-                                                                   optimizer=data_generation_config.optimizer,
-                                                                   image_pipeline=image_pipeline,
-                                                                   activation_extractor=activation_extractor,
-                                                                   image_granularity=data_generation_config.image_granularity,
-                                                                   scheduler_step_fn=scheduler_step_fn,
-                                                                   scheduler=scheduler,
-                                                                   initial_lr=data_generation_config.initial_lr,
-                                                                   normalization_mean=normalization[0],
-                                                                   normalization_std=normalization[1],
-                                                                   clip_images=data_generation_config.clip_images,
-                                                                   reflection=data_generation_config.reflection)
+                                                                data_gen_batch_size=data_generation_config.data_gen_batch_size,
+                                                                init_dataset=init_dataset,
+                                                                optimizer=data_generation_config.optimizer,
+                                                                image_pipeline=image_pipeline,
+                                                                activation_extractor=activation_extractor,
+                                                                image_granularity=data_generation_config.image_granularity,
+                                                                scheduler_step_fn=scheduler_step_fn,
+                                                                scheduler=scheduler,
+                                                                initial_lr=data_generation_config.initial_lr,
+                                                                normalization_mean=normalization[0],
+                                                                normalization_std=normalization[1],
+                                                                device=device)
 
         # Perform data generation and obtain a list of generated images
         generated_images_list = data_generation(
@@ -247,6 +249,7 @@ if FOUND_TORCH and FOUND_TORCHVISION:
             bn_alignment_loss_fn=bn_alignment_loss_fn,
             output_loss_fn=output_loss_fn,
             output_loss_multiplier=data_generation_config.output_loss_multiplier,
+            device=device,
         )
         # Return the list of finalized generated images
         return generated_images_list
@@ -261,7 +264,8 @@ if FOUND_TORCH and FOUND_TORCHVISION:
             bn_layer_weighting_fn: Callable,
             bn_alignment_loss_fn: Callable,
             output_loss_fn: Callable,
-            output_loss_multiplier: float
+            output_loss_multiplier: float,
+            device: torch.device
     ) -> List[Any]:
         """
         Function to perform data generation using the provided model and data generation configuration.
@@ -276,14 +280,11 @@ if FOUND_TORCH and FOUND_TORCHVISION:
             bn_alignment_loss_fn (Callable): Function to compute BatchNorm alignment loss.
             output_loss_fn (Callable): Function to compute output loss.
             output_loss_multiplier (float): Multiplier for the output loss.
+            device (torch.device): The current device set for PyTorch operations.
 
         Returns:
             List: Finalized list containing generated images.
         """
-
-        # Compute the layer weights based on orig_bn_stats_holder
-        bn_layer_weights = bn_layer_weighting_fn(orig_bn_stats_holder)
-
         # Get the current time to measure the total time taken
         total_time = time.time()
 
@@ -291,7 +292,7 @@ if FOUND_TORCH and FOUND_TORCHVISION:
         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()
@@ -311,7 +312,11 @@ if FOUND_TORCH and FOUND_TORCHVISION:
                 input_imgs = image_pipeline.image_input_manipulation(imgs_to_optimize)
 
                 # Forward pass to extract activations
-                output = activation_extractor.run_model(input_imgs)
+                with autocast():
+                    output = activation_extractor.run_model(input_imgs)
+
+                # Compute the layer weights based on orig_bn_stats_holder
+                bn_layer_weights = bn_layer_weighting_fn(orig_bn_stats_holder, activation_extractor, i_iter, data_generation_config.n_iter)
 
                 # Compute BatchNorm alignment loss
                 bn_loss = all_imgs_opt_handler.compute_bn_loss(input_imgs=input_imgs,
@@ -322,33 +327,33 @@ if FOUND_TORCH and FOUND_TORCHVISION:
                                                                bn_layer_weights=bn_layer_weights)
 
                 # Compute output loss
-                if output_loss_multiplier > 0:
-                    output_loss = output_loss_fn(
-                        output_imgs=output,
-                        activation_extractor=activation_extractor)
-                else:
-                    output_loss = torch.zeros(1).to(get_working_device())
+                output_loss = output_loss_fn(
+                    model_outputs=output,
+                    activation_extractor=activation_extractor,
+                    device=device)
 
                 # Compute total loss
                 total_loss = bn_loss + output_loss_multiplier * output_loss
 
                 # Perform optimiztion step
-                all_imgs_opt_handler.optimization_step(random_batch_index, total_loss, i_ter)
+                all_imgs_opt_handler.optimization_step(random_batch_index, total_loss, i_iter)
 
                 # Update the statistics based on the updated images
                 if all_imgs_opt_handler.use_all_data_stats:
-                    final_imgs = image_pipeline.image_output_finalize(imgs_to_optimize)
-                    all_imgs_opt_handler.update_statistics(input_imgs=final_imgs,
-                                                           batch_index=random_batch_index,
-                                                           activation_extractor=activation_extractor)
+                    with autocast():
+                        final_imgs = image_pipeline.image_output_finalize(imgs_to_optimize)
+                        all_imgs_opt_handler.update_statistics(input_imgs=final_imgs,
+                                                               batch_index=random_batch_index,
+                                                               activation_extractor=activation_extractor)
 
             ibar.set_description(f"Total Loss: {total_loss.item():.5f}, "
                                  f"BN Loss: {bn_loss.item():.5f}, "
-                                 f"Output Loss: {output_loss_multiplier * output_loss.item():.5f}")
+                                 f"Output Loss: {output_loss.item():.5f}")
 
         # Return a list containing the finalized generated images
         finalized_imgs = all_imgs_opt_handler.get_finalized_images()
         Logger.info(f'Total time to generate {len(finalized_imgs)} images (seconds): {int(time.time() - total_time)}')
+        Logger.info(f'Final Loss: Total {total_loss.item()}, BN loss {bn_loss.item()}, Output loss {output_loss.item()}')
         return finalized_imgs
 else:
     # If torch is not installed,