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

model_compression_toolkit/data_generation/keras/optimization_functions/lr_scheduler.py

@@ -0,0 +1,219 @@
+# 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.
+# ==============================================================================
+from typing import Dict, Union
+
+import numpy as np
+import tensorflow as tf
+
+from model_compression_toolkit.logger import Logger
+
+
+class ReduceLROnPlateau(tf.keras.callbacks.Callback):
+    """
+    Reduce learning rate when a metric has stopped improving.
+    """
+
+    def __init__(self, optimizer: tf.keras.optimizers.Optimizer, mode: str = 'min', factor: float = 0.1,
+                 patience: int = 10, threshold: float = 1e-4, threshold_mode: str = 'rel', cooldown: int = 0,
+                 min_lr: float = 0, eps: float = 1e-8, verbose: bool = False):
+        """
+       Initialize the ReduceLROnPlateau scheduler.
+
+       Args:
+           optimizer (tf.keras.optimizers.Optimizer): Wrapped optimizer.
+           mode (str): One of `min`, `max`. In `min` mode, lr will be reduced when the quantity
+                       monitored has stopped decreasing; in `max` mode it will be reduced when the
+                       quantity monitored has stopped increasing. Default: 'min'.
+           factor (float): Factor by which the learning rate will be reduced. new_lr = lr * factor.
+                           Default: 0.1.
+           patience (int): Number of epochs with no improvement after which learning rate will be reduced.
+                           Default: 10.
+           threshold (float): Threshold for measuring the new optimum, to only focus on significant changes.
+                              Default: 1e-4.
+           threshold_mode (str): One of `rel`, `abs`. In `rel` mode, dynamic_threshold = best * ( 1 + threshold )
+                                 in 'max' mode or best * ( 1 - threshold ) in `min` mode. In `abs` mode, dynamic_threshold
+                                 = best + threshold in `max` mode or best - threshold in `min` mode. Default: 'rel'.
+           cooldown (int): Number of epochs to wait before resuming normal operation after lr has been reduced.
+                           Default: 0.
+           min_lr (float): A lower bound on the learning rate. Default: 0.
+           eps (float): Minimal decay applied to lr. If the difference between new and old lr is smaller than eps,
+                        the update is ignored. Default: 1e-8.
+           verbose (bool): If True, prints a message to stdout for each update. Default: False.
+       """
+
+        super(ReduceLROnPlateau, self).__init__()
+
+        if factor >= 1.0:
+            Logger.critical('Factor should be < 1.0.') # pragma: no cover
+        self.factor = factor
+
+        self.optimizer = optimizer
+        self.min_lr = min_lr
+        self.patience = patience
+        self.verbose = verbose
+        self.cooldown = cooldown
+        self.cooldown_counter = 0
+        self.mode = mode
+        self.threshold = threshold
+        self.threshold_mode = threshold_mode
+        self.best = None
+        self.num_bad_epochs = None
+        self.mode_worse = None  # the worse value for the chosen mode
+        self.eps = eps
+        self.last_epoch = 0
+        self._init_is_better(mode=mode, threshold=threshold, threshold_mode=threshold_mode)
+        self._reset()
+
+    def _reset(self) -> None:
+        """
+        Resets num_bad_epochs counter and cooldown counter.
+        """
+        self.best = self.mode_worse
+        self.cooldown_counter = 0
+        self.num_bad_epochs = 0
+
+    def on_epoch_end(self, epoch: int, loss: float, logs: dict = None) -> None:
+        """
+        Check conditions and update learning rate at the end of an epoch.
+
+        Args:
+            epoch (int): The current epoch number.
+            loss (float): Validation loss value.
+            logs (dict): The dictionary of logs from the epoch.
+        """
+        current = float(loss)
+
+        if self.is_better(current, self.best):
+            self.best = current
+            self.num_bad_epochs = 0
+        else:
+            self.num_bad_epochs += 1
+
+        if self.in_cooldown:
+            self.cooldown_counter -= 1
+            self.num_bad_epochs = 0  # Ignore any bad epochs in cooldown
+
+        if self.num_bad_epochs > self.patience:
+            self._reduce_lr(epoch)
+            self.cooldown_counter = self.cooldown
+            self.num_bad_epochs = 0
+            self.best = self.mode_worse
+
+    def _reduce_lr(self, epoch: int) -> None:
+        """
+       Reduce the learning rate for each parameter group.
+
+       Args:
+           epoch (int): The current epoch number.
+       """
+        old_lr = float(tf.keras.backend.get_value(self.optimizer.learning_rate))
+        new_lr = max(old_lr * self.factor, self.min_lr)
+        if old_lr - new_lr > self.eps:
+            tf.keras.backend.set_value(self.optimizer.learning_rate, new_lr)
+            if self.verbose:
+                print(f'Epoch {epoch:05d}: reducing learning rate to {new_lr:.4e}.')
+
+    @property
+    def in_cooldown(self) -> bool:
+        """
+        Check if the scheduler is in a cooldown period.
+
+        Returns:
+            bool: True if in cooldown period, False otherwise.
+        """
+        return self.cooldown_counter > 0
+
+    def is_better(self, a: float, best: Union[float, None]) -> bool:
+        """
+        Determine if the new value is better than the best value based on mode and threshold.
+
+        Args:
+            a (float): The new value to compare.
+            best (float): The best value to compare against.
+
+        Returns:
+            bool: True if the new value is better, False otherwise.
+        """
+        if best is None:
+            return True
+
+        if self.mode == 'min' and self.threshold_mode == 'rel':
+            rel_epsilon = 1. - self.threshold
+            return a < best * rel_epsilon
+        elif self.mode == 'min' and self.threshold_mode == 'abs':
+            return a < best - self.threshold
+        elif self.mode == 'max' and self.threshold_mode == 'rel':
+            rel_epsilon = self.threshold + 1.
+            return a > best * rel_epsilon
+        else:  # mode == 'max' and threshold_mode == 'abs':
+            return a > best + self.threshold
+
+    def _init_is_better(self, mode: str, threshold: float, threshold_mode: str) -> None:
+        """
+        Initialize the comparison function for determining if a new value is better.
+
+        Args:
+            mode (str): The mode for comparison, 'min' or 'max'.
+            threshold (float): The threshold for comparison.
+            threshold_mode (str): The mode for threshold, 'rel' or 'abs'.
+
+        Raises:
+            ValueError: If an unknown mode or threshold mode is provided.
+        """
+        if mode not in {'min', 'max'}:
+            Logger.critical(f'mode {mode} is unknown!') # pragma: no cover
+        if threshold_mode not in {'rel', 'abs'}:
+            Logger.critical(f'threshold mode {threshold_mode} is unknown!') # pragma: no cover
+
+        if mode == 'min':
+            self.mode_worse = float('inf')
+        else:  # mode == 'max':
+            self.mode_worse = float('-inf')
+
+        self.mode = mode
+        self.threshold = threshold
+        self.threshold_mode = threshold_mode
+
+    def get_config(self) -> Dict:
+        """
+        Return the configuration of the scheduler as a dictionary.
+
+        Returns:
+            Dict: The configuration of the scheduler.
+        """
+        config = {
+            'factor': self.factor,
+            'patience': self.patience,
+            'verbose': self.verbose,
+            'mode': self.mode,
+            'threshold': self.threshold,
+            'threshold_mode': self.threshold_mode,
+            'cooldown': self.cooldown,
+            'min_lr': self.min_lr,
+            'eps': self.eps
+        }
+        base_config = super(ReduceLROnPlateau, self).get_config()
+        return {**base_config, **config}
+
+    def set_config(self, config: Dict) -> None:
+        """
+        Set the configuration of the scheduler from a dictionary.
+
+        Args:
+            config (Dict): The configuration dictionary.
+        """
+        for key, value in config.items():
+            setattr(self, key, value)
+

model_compression_toolkit/data_generation/keras/optimization_functions/output_loss_functions.py

@@ -22,7 +22,7 @@ from model_compression_toolkit.data_generation.keras.model_info_exctractors impo
 
 # Function to calculate the regularized min-max difference loss
 def regularized_min_max_diff(
-        output_imgs: tf.Tensor,
+        model_outputs: tf.Tensor,
         activation_extractor: KerasActivationExtractor,
         tape: tf.GradientTape,
         eps: float = 1e-6,
@@ -33,7 +33,7 @@ def regularized_min_max_diff(
     This function calculates the regularized min-max difference loss based on the provided inputs.
 
     Args:
-        output_imgs (tf.Tensor): Output images or tensors.
+        model_outputs (tf.Tensor): Output images or tensors.
         activation_extractor (KerasActivationExtractor): Activation extractor object.
         tape (tf.GradientTape): TensorFlow tape for recording operations.
         eps (float, optional): Small constant to prevent division by zero.
@@ -82,38 +82,63 @@ def regularized_min_max_diff(
     return output_loss
 
 
-def min_max_diff(
-        output_imgs: tf.Tensor,
+def inverse_min_max_diff(
+        model_outputs: tf.Tensor,
         eps: float = 1e-6,
         **kwargs) -> tf.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.
         eps (float): Small value for numerical stability.
         **kwargs: Additional keyword arguments.
 
     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 = tf.zeros(1)
+    for output in model_outputs:
         output = tf.reshape(output, [output.shape[0], -1])
         output_loss += 1 / (tf.reduce_max(output, 1) - tf.reduce_min(output, 1) + eps)
     return output_loss
 
+def negative_min_max_diff(
+        model_outputs: tf.Tensor,
+        eps: float = 1e-6,
+        **kwargs) -> tf.Tensor:
+    """
+    Calculate the inverse of the 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.
+        eps (float): Small value for numerical stability.
+        **kwargs: Additional keyword arguments.
+
+    Returns:
+        Tensor: The computed minimum-maximum difference loss.
+    """
+    if not isinstance(model_outputs, (list, tuple)):
+        model_outputs = [model_outputs]
+    output_loss = tf.zeros(1)
+    for output in model_outputs:
+        output = tf.reshape(output, [output.shape[0], -1])
+        out_max = tf.reduce_max(output, 1)
+        out_min = tf.reduce_min(output, 1)
+        output_loss += tf.reduce_mean(-(out_max - out_min))
+    return output_loss
+
 
 def no_output_loss(
-        output_imgs: tf.Tensor,
+        model_outputs: tf.Tensor,
         **kwargs) -> tf.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.
         **kwargs: Additional keyword arguments.
 
     Returns:
@@ -125,6 +150,7 @@ def no_output_loss(
 # 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/keras/optimization_functions/scheduler_step_functions.py

@@ -18,104 +18,11 @@ from typing import Callable, Dict
 import numpy as np
 
 from model_compression_toolkit.data_generation.common.enums import SchedulerType
+from model_compression_toolkit.data_generation.keras.optimization_functions.lr_scheduler import \
+    ReduceLROnPlateau
 
 
-# Custom implementation of Reduce LR On Plateau schedular
-# Customized for gradient taping
-class CustomReduceLROnPlateau:
-    def __init__(self,
-                 factor: float = 0.5,
-                 patience: int = 10,
-                 min_delta: float = 1e-4,
-                 cooldown: int = 0,
-                 min_lr: float = 1e-6,
-                 sign_number: int = 4,
-                 optim_lr=None,
-                 ):
-        """
-        Initialize a custom learning rate scheduler based on ReduceLROnPlateau.
-
-        Args:
-            factor (float): Factor by which the learning rate will be reduced.
-            patience (int): Number of epochs with no improvement after which learning rate will be reduced.
-            min_delta (float): Minimum change in monitored value to qualify as an improvement.
-            cooldown (int): Number of epochs to wait before resuming after reducing the learning rate.
-            min_lr (float): Lower bound on the learning rate.
-            sign_number (int): Number of significant digits to consider for comparisons when checking for improvement.
-            optim_lr (tf.Variable): Optimizer learning rate variable to synchronize with the reduced learning rate.
-        """
-
-        self.optim_lr = optim_lr
-        self.factor = factor
-        self.min_lr = min_lr
-        self.min_delta = min_delta
-        self.patience = patience
-        self.cooldown = cooldown
-        self.cooldown_counter = 0
-        self.wait = 0
-        self.best = 0
-        self.monitor_op = None
-        self.sign_number = sign_number
-        self.reduce_lr = True
-        self._reset()
-
-    def _reset(self):
-        """
-        Reset the internal state of the learning rate scheduler.
-        """
-        self.monitor_op = lambda a, b: np.less(a, b - self.min_delta)
-        self.best = np.Inf
-        self.cooldown_counter = 0
-        self.wait = 0
-
-    def on_epoch_end(self,
-                     loss: float,
-                     logs=None):
-        """
-        Update the learning rate based on the validation loss at the end of each epoch.
-
-        Args:
-            loss (float): Validation loss value.
-            logs (dict): Dictionary of training metrics and logs.
-
-        Notes:
-            This method should be called at the end of each epoch during training.
-        """
-        logs = logs or {}
-        logs['lr'] = float(self.optim_lr.learning_rate.numpy())
-        current = float(loss)
-
-        if self.in_cooldown():
-            self.cooldown_counter -= 1
-            self.wait = 0
-
-        if self.monitor_op(current, self.best):
-            self.best = current
-            self.wait = 0
-        elif not self.in_cooldown():
-            self.wait += 1
-            if self.wait >= self.patience:
-
-                old_lr = float(self.optim_lr.learning_rate.numpy())
-                if old_lr > self.min_lr and self.reduce_lr:
-                    new_lr = old_lr * self.factor
-
-                    new_lr = max(new_lr, self.min_lr)
-                    self.optim_lr.learning_rate.assign(new_lr)
-                    self.cooldown_counter = self.cooldown
-                    self.wait = 0
-
-    def in_cooldown(self) -> bool:
-        """
-        Check if the learning rate scheduler is in the cooldown phase.
-
-        Returns:
-            bool: True if in cooldown, False otherwise.
-        """
-        return self.cooldown_counter > 0
-
-
-def get_reduceonplatue_scheduler(n_iter: int, initial_lr: float):
+def get_reduce_lr_on_plateau_scheduler(n_iter: int, initial_lr: float):
     """
     Create a custom ReduceLROnPlateau learning rate scheduler.
 
@@ -126,10 +33,11 @@ def get_reduceonplatue_scheduler(n_iter: int, initial_lr: float):
     Returns:
         callable: Partial function for creating CustomReduceLROnPlateau scheduler.
     """
-    return partial(CustomReduceLROnPlateau)
+    return partial(ReduceLROnPlateau, min_lr=1e-4, factor=0.5, patience=int(n_iter / 50))
+
 
 
 # Define a dictionary that maps scheduler types to functions for creating schedulers.
 scheduler_step_function_dict: Dict[SchedulerType, Callable] = {
-    SchedulerType.REDUCE_ON_PLATEAU: get_reduceonplatue_scheduler  # Custom ReduceLROnPlateau scheduler.
+    SchedulerType.REDUCE_ON_PLATEAU: get_reduce_lr_on_plateau_scheduler,  # ReduceLROnPlateau scheduler.
 }

model_compression_toolkit/data_generation/keras/optimization_utils.py

@@ -25,6 +25,7 @@ from model_compression_toolkit.data_generation.common.optimization_utils import
     BatchOptimizationHolder, AllImagesStatsHolder, BatchStatsHolder
 from model_compression_toolkit.data_generation.keras.constants import IMAGE_MIN_VAL, IMAGE_MAX_VAL, BATCH_AXIS, \
     H_AXIS, W_AXIS
+from model_compression_toolkit.data_generation.keras.image_operations import create_valid_grid
 from model_compression_toolkit.data_generation.keras.model_info_exctractors import KerasActivationExtractor, \
     KerasOriginalBNStatsHolder
 
@@ -64,11 +65,7 @@ class KerasImagesOptimizationHandler(ImagesOptimizationHandler):
         # IMAGE_MIN_VAL( default set to 0) - IMAGE_MAX_VAL ( default set to 255) before normalization
         self.normalization_mean = normalization_mean
         self.normalization_std = normalization_std
-        self.valid_grid = []
-        for i, (mean, var) in enumerate(zip(self.normalization_mean, self.normalization_std)):
-            min_val = (IMAGE_MIN_VAL - mean) / var
-            max_val = (IMAGE_MAX_VAL - mean) / var
-            self.valid_grid.append((min_val, max_val))
+        self.valid_grid = create_valid_grid(self.normalization_mean, self.normalization_std)
 
         super(KerasImagesOptimizationHandler, self).__init__(model=model,
                                                              data_gen_batch_size=
@@ -83,8 +80,6 @@ class KerasImagesOptimizationHandler(ImagesOptimizationHandler):
                                                              initial_lr=data_generation_config.initial_lr,
                                                              normalization_mean=self.normalization_mean,
                                                              normalization_std=self.normalization_std,
-                                                             clip_images=data_generation_config.clip_images,
-                                                             reflection=data_generation_config.reflection,
                                                              eps=eps)
 
         # Set the mean axis based on the image granularity
@@ -140,22 +135,14 @@ class KerasImagesOptimizationHandler(ImagesOptimizationHandler):
         Returns:
             tf.Tensor: Clipped and reflected tensor.
         """
-        if self.clip_images:
-            images = z.numpy()
-            for i_ch in range(len(self.valid_grid)):
-                # Clip the values of the channel within the valid range.
-                clamp = tf.clip_by_value(t=z[:, :, :, i_ch], clip_value_min=self.valid_grid[i_ch][0],
-                                         clip_value_max=self.valid_grid[i_ch][1])
-                if self.reflection:
-                    # Reflect the values.
-                    images[:, :, :, i_ch] = 2 * clamp - z[:, :, :, i_ch]
-                else:
-                    images[:, :, :, i_ch] = clamp
-            # Assign the clipped reflected values back to `z`.
-            z.assign(images)
-            return z
-        else:
-            return z
+        images = z.numpy()
+        for i_ch in range(len(self.valid_grid)):
+            # Clip the values of the channel within the valid range.
+            clamp = tf.clip_by_value(t=z[:, :, :, i_ch], clip_value_min=self.valid_grid[i_ch][0],
+                                     clip_value_max=self.valid_grid[i_ch][1])
+        # Assign the clipped reflected values back to `z`.
+        z.assign(images)
+        return z
 
     def get_layer_accumulated_stats(self, layer_name: str) -> Tuple[tf.Tensor, tf.Tensor]:
         """
@@ -255,7 +242,7 @@ class KerasImagesOptimizationHandler(ImagesOptimizationHandler):
                           images: tf.Tensor,
                           gradients: tf.Tensor,
                           loss: tf.Tensor,
-                          i_ter: int):
+                          i_iter: int):
         """
         Perform an optimization step.
 
@@ -264,7 +251,7 @@ class KerasImagesOptimizationHandler(ImagesOptimizationHandler):
             images (tf.Tensor): The images to optimize for the batch.
             gradients (List[tf.Tensor]): The gradients calculated for the images.
             loss (tf.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=batch_index)
@@ -274,7 +261,7 @@ class KerasImagesOptimizationHandler(ImagesOptimizationHandler):
         optimizer.apply_gradients(zip(gradients, [images]))
 
         # Perform scheduler step
-        scheduler.on_epoch_end(loss=tf.reduce_mean(loss))
+        scheduler.on_epoch_end(epoch=i_iter, loss=tf.reduce_mean(loss))
 
     def get_finilized_data_loader(self) -> np.ndarray:
         """
@@ -321,8 +308,8 @@ class KerasBatchOptimizationHolder(BatchOptimizationHolder):
             initial_lr (float): The initial learning rate used by the optimizer.
         """
         self.images = images
-        self.optimizer = optimizer(lr=initial_lr)
-        self.scheduler = scheduler(optim_lr=self.optimizer)
+        self.optimizer = optimizer(learning_rate=initial_lr)
+        self.scheduler = scheduler(optimizer=self.optimizer)
 
 
 class KerasAllImagesStatsHolder(AllImagesStatsHolder):

model_compression_toolkit/data_generation/pytorch/constants.py

@@ -20,8 +20,11 @@ BATCH_AXIS, CHANNEL_AXIS, H_AXIS, W_AXIS = 0, 1, 2, 3
 # Default initial learning rate constant.
 DEFAULT_PYTORCH_INITIAL_LR = 16
 
+# Default extra pixels for image padding.
+DEFAULT_PYTORCH_EXTRA_PIXELS = 32
+
 # Default output loss multiplier.
-DEFAULT_PYTORCH_OUTPUT_LOSS_MULTIPLIER = 1e-6
+DEFAULT_PYTORCH_OUTPUT_LOSS_MULTIPLIER = 1e-5
 
 # Default BatchNorm layer types
 DEFAULT_PYTORCH_BN_LAYER_TYPES = [torch.nn.BatchNorm2d]

model_compression_toolkit/data_generation/pytorch/image_operations.py

@@ -0,0 +1,105 @@
+# 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.
+# ==============================================================================
+from typing import List
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from model_compression_toolkit.core.pytorch.pytorch_device_config import get_working_device
+from torchvision.transforms import Normalize
+
+from model_compression_toolkit.logger import Logger
+
+
+def create_valid_grid(means: List[int], stds: List[int]) -> torch.Tensor:
+    """
+    Create a valid grid for image normalization.
+
+    Args:
+        means (List[int]): List of mean values per channel.
+        stds (List[int]): List of standard deviation values per channel.
+
+    Returns:
+        torch.Tensor: The valid grid for image normalization.
+    """
+    # Image valid grid
+    pixel_grid = torch.from_numpy(np.array(list(range(256))).repeat(3).reshape(-1, 3)).float()
+    valid_grid = Normalize(mean=means, std=stds)(pixel_grid.transpose(1, 0)[None, :, :, None]).squeeze().to(
+        get_working_device())
+    return valid_grid
+
+
+class Smoothing(torch.nn.Module):
+    """
+    A PyTorch module for applying Gaussian smoothing to an image.
+    """
+
+    def __init__(self, size: int = 3, sigma: float = 1.25, kernel: torch.Tensor = None):
+        """
+        Initialize the Smoothing module.
+
+        Args:
+            size (int): The size of the Gaussian kernel (Default: 3).
+            sigma (float): The standard deviation of the Gaussian kernel (Defalut: 1.25).
+            kernel (torch.Tensor, optional): Precomputed Gaussian kernel. If None, it will be created.
+        """
+        super().__init__()
+        if kernel is None:
+            kernel = self.gaussian_kernel(size, sigma)
+        if kernel.dim() != 2:
+            Logger.critical("Kernel must have 2 dimensions. Found {} dimensions.".format(kernel.dim())) # pragma: no cover
+        kernel = kernel.view(1, 1, kernel.shape[0], kernel.shape[1])
+        # Repeat for 3 color channels
+        kernel = kernel.repeat(3, 1, 1, 1)
+        self.kernel = kernel
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        """
+        Apply Gaussian smoothing to the input image.
+
+        Args:
+            image (torch.Tensor): The input image tensor.
+
+        Returns:
+            torch.Tensor: The smoothed image tensor.
+        """
+        return F.conv2d(image, self.kernel.to(image.device), padding=self.kernel.shape[-1] // 2, groups=3)
+
+    def __repr__(self) -> str:
+        """
+        Return the string representation of the Smoothing module.
+
+        Returns:
+            str: String representation of the Smoothing module.
+        """
+        return f"{self.__class__.__name__}(kernel={self.kernel.shape[-1]})" # pragma: no cover
+
+    @staticmethod
+    def gaussian_kernel(size: int = 3, sigma: float = 1) -> torch.Tensor:
+        """
+        Create a Gaussian kernel.
+
+        Args:
+            size (int): The size of the Gaussian kernel.
+            sigma (float): The standard deviation of the Gaussian kernel.
+
+        Returns:
+            torch.Tensor: The Gaussian kernel tensor.
+        """
+        axis = torch.arange(-size // 2 + 1., size // 2 + 1.)
+        x, y = torch.meshgrid(axis, axis)
+        kernel = torch.exp(-(x ** 2 + y ** 2) / (2 * sigma ** 2))
+        kernel = kernel / torch.sum(kernel)
+        return kernel