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/image_pipeline.py

@@ -12,36 +12,45 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Type, Dict, List
+import numpy as np
+import torch
+from typing import Type, Dict, Tuple, Union, List
 
 from torch import Tensor
 from torchvision.transforms import RandomCrop, RandomHorizontalFlip, CenterCrop
 
-
-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.pytorch.image_operations import Smoothing, create_valid_grid
 
 
 class PytorchIdentityImagePipeline(BaseImagePipeline):
     """
     An image pipeline implementation for PyTorch models that returns the input images as is (identity).
     """
-    def __init__(self, output_image_size: int, extra_pixels: int = 0):
+    def __init__(self,
+                 output_image_size: Union[int, Tuple[int, int]],
+                 extra_pixels: Union[int, Tuple[int, int]] = 0,
+                 normalization: List[List[int]] = [[0, 0, 0], [1, 1, 1]],
+                 image_clipping: bool = True,
+                 ):
         """
         Initialize the PytorchIdentityImagePipeline.
 
         Args:
-            output_image_size (int): The output image size.
-            extra_pixels (int): Extra pixels to add to the input image size (not used in identity pipeline).
+            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 (not used in identity pipeline).
+            normalization (List[List[float]]): The image normalization values for processing images during optimization.
+            image_clipping (bool): Whether to clip images during optimization.
         """
-        super(PytorchIdentityImagePipeline, self).__init__(output_image_size)
+        super(PytorchIdentityImagePipeline, self).__init__(output_image_size, extra_pixels, image_clipping, normalization)
 
-    def get_image_input_size(self) -> int:
+    def get_image_input_size(self) -> Tuple[int, int]:
         """
         Get the input size of the image.
 
         Returns:
-            int: The input image size.
+            Tuple[int, int]: The input image size.
         """
         return self.output_image_size
 
@@ -70,31 +79,43 @@ class PytorchIdentityImagePipeline(BaseImagePipeline):
         return images
 
 
-class PytorchRandomCropImagePipeline(BaseImagePipeline):
+class PytorchSmoothAugmentationImagePipeline(BaseImagePipeline):
     """
-    An image pipeline implementation for PyTorch models that includes random cropping.
+    An image pipeline implementation for PyTorch models that includes random cropping and flipping.
     """
-    def __init__(self, output_image_size: int, extra_pixels: int = 0):
+    def __init__(self,
+                 output_image_size: Union[int, Tuple[int, int]],
+                 extra_pixels: Union[int, Tuple[int, int]] = 0,
+                 normalization: List[List[int]] = [[0, 0, 0], [1, 1, 1]],
+                 image_clipping: bool = True,
+                 smoothing_filter_size: int = 3,
+                 smoothing_filter_sigma: float = 1.25):
         """
         Initialize the PytorchRandomCropFlipImagePipeline.
 
         Args:
-            output_image_size (int): The output image size.
-            extra_pixels (int): Extra pixels to add to the input image size. Defaults to 0.
-        """
-        super(PytorchRandomCropImagePipeline, self).__init__(output_image_size)
-        self.extra_pixels = 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(PytorchSmoothAugmentationImagePipeline, self).__init__(output_image_size, extra_pixels, image_clipping, normalization)
+        self.smoothing = Smoothing(size=smoothing_filter_size, sigma=smoothing_filter_sigma)
         self.random_crop = RandomCrop(self.output_image_size)
+        self.random_flip = RandomHorizontalFlip(0.5)
         self.center_crop = CenterCrop(self.output_image_size)
+        self.valid_grid = create_valid_grid(means=self.normalization[0], stds=self.normalization[1])
 
-    def get_image_input_size(self) -> int:
+    def get_image_input_size(self) -> Tuple[int, int]:
         """
         Get the input size of the image.
 
         Returns:
-            int: The input image size.
+            Tuple[int, int]: The input image size.
         """
-        return 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: Tensor) -> Tensor:
         """
@@ -104,9 +125,14 @@ class PytorchRandomCropImagePipeline(BaseImagePipeline):
             images (Tensor): The input images.
 
         Returns:
-            Tensor: The manipulated images (randomly flipped and cropped).
+            Tensor: The manipulated images.
         """
-        return self.random_crop(images)
+        new_images = self.random_flip(images)
+        new_images = self.smoothing(new_images)
+        new_images = self.random_crop(new_images)
+        if self.image_clipping:
+            new_images = self.clip_images(new_images, self.valid_grid)
+        return new_images
 
     def image_output_finalize(self, images: Tensor) -> Tensor:
         """
@@ -118,71 +144,37 @@ class PytorchRandomCropImagePipeline(BaseImagePipeline):
         Returns:
             Tensor: The finalized images (center cropped).
         """
-        return self.center_crop(images)
-
+        new_images = self.smoothing(images)
+        new_images = self.center_crop(new_images)
+        if self.image_clipping:
+            new_images = self.clip_images(new_images, self.valid_grid)
+        return new_images
 
-class PytorchRandomCropFlipImagePipeline(BaseImagePipeline):
-    """
-    An image pipeline implementation for PyTorch models that includes random cropping and flipping.
-    """
-    def __init__(self, output_image_size: int, extra_pixels: int = 0):
+    @staticmethod
+    def clip_images(images: Tensor, valid_grid: Tensor, reflection: bool = False) -> Tensor:
         """
-        Initialize the PytorchRandomCropFlipImagePipeline.
+        Clip the images based on a valid grid.
 
         Args:
-            output_image_size (int): The output image size.
-            extra_pixels (int): Extra pixels to add to the input image size. Defaults to 0.
-        """
-        super(PytorchRandomCropFlipImagePipeline, self).__init__(output_image_size)
-        self.extra_pixels = extra_pixels
-        self.random_crop = RandomCrop(self.output_image_size)
-        self.random_flip = RandomHorizontalFlip(0.5)
-        self.center_crop = CenterCrop(self.output_image_size)
-
-    def get_image_input_size(self) -> int:
-        """
-        Get the input size of the image.
+            images (Tensor): The images to be clipped.
+            valid_grid (Tensor): The valid grid for clipping.
+            reflection (bool): Whether to apply reflection during clipping. Defaults to False.
 
         Returns:
-            int: The input image size.
-        """
-        return self.output_image_size + self.extra_pixels
-
-    def image_input_manipulation(self, images: Tensor) -> Tensor:
-        """
-        Manipulate the input images with random flipping and cropping.
-
-        Args:
-            images (Tensor): The input images.
-
-        Returns:
-            Tensor: The manipulated images (randomly flipped and cropped).
-        """
-        random_flipped_data = self.random_flip(images)
-        return self.random_crop(random_flipped_data)
-
-    def image_output_finalize(self, images: Tensor) -> Tensor:
-        """
-        Finalize the output images with center cropping.
-
-        Args:
-            images (Tensor): The output images.
-
-        Returns:
-            Tensor: The finalized images (center cropped).
-        """
-        return self.center_crop(images)
+            Tensor: The clipped images.
+        """
+        with torch.no_grad():
+            for i_ch in range(valid_grid.shape[0]):
+                clamp = torch.clamp(images[:, i_ch, :, :], valid_grid[i_ch, :].min(), valid_grid[i_ch, :].max())
+                if reflection:
+                    images[:, i_ch, :, :] = 2 * clamp - images[:, i_ch, :, :]
+                else:
+                    images[:, i_ch, :, :] = clamp
+        return images
 
 
 # Dictionary mapping ImagePipelineType to corresponding image pipeline classes
 image_pipeline_dict: Dict[ImagePipelineType, Type[BaseImagePipeline]] = {
     ImagePipelineType.IDENTITY: PytorchIdentityImagePipeline,
-    ImagePipelineType.RANDOM_CROP: PytorchRandomCropImagePipeline,
-    ImagePipelineType.RANDOM_CROP_FLIP: PytorchRandomCropFlipImagePipeline
-}
-
-# 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.NO_NORMALIZATION: [[0, 0, 0], [1, 1, 1]]
+    ImagePipelineType.SMOOTHING_AND_AUGMENTATION: PytorchSmoothAugmentationImagePipeline
 }

model_compression_toolkit/data_generation/pytorch/model_info_exctractors.py

@@ -126,7 +126,6 @@ class PytorchActivationExtractor(ActivationExtractor):
         self.layer_types_to_extract_inputs = tuple(layer_types_to_extract_inputs)
         self.last_layer_types_to_extract_inputs = tuple(last_layer_types_to_extract_inputs)
         self.num_layers = sum([1 if isinstance(layer, tuple(layer_types_to_extract_inputs)) else 0 for layer in model.modules()])
-        Logger.info(f'Number of layers = {self.num_layers}')
         self.hooks = {}  # Dictionary to store InputHook instances by layer name
         self.last_linear_layers_hooks = {}  # Dictionary to store InputHook instances by layer name
         self.hook_handles = []  # List to store hook handles
@@ -206,15 +205,6 @@ class PytorchActivationExtractor(ActivationExtractor):
         """
         return self.last_linear_layer_weights
 
-    def get_num_extractor_layers(self) -> int:
-        """
-        Get the number of hooked layers in the model.
-
-        Returns:
-            int: Number of hooked layers in the model.
-        """
-        return self.num_layers
-
     def get_extractor_layer_names(self) -> List:
         """
         Get a list of the hooked layer names.

model_compression_toolkit/data_generation/pytorch/optimization_functions/bn_layer_weighting_functions.py

@@ -18,16 +18,21 @@ import torch
 
 from model_compression_toolkit.data_generation.common.enums import BNLayerWeightingType
 from model_compression_toolkit.data_generation.pytorch.model_info_exctractors import OriginalBNStatsHolder, \
-    ActivationExtractor
+    ActivationExtractor, PytorchActivationExtractor
 
 
-def average_bn_layer_weighting_fn(orig_bn_stats_holder: OriginalBNStatsHolder, **kwargs) -> Dict[str, float]:
+def average_bn_layer_weighting_fn(orig_bn_stats_holder: OriginalBNStatsHolder,
+                                  activation_extractor: PytorchActivationExtractor,
+                                  i_iter: int,
+                                  n_iter: int) -> Dict[str, float]:
     """
     Calculate average weighting for each batch normalization layer.
 
     Args:
         orig_bn_stats_holder (OriginalBNStatsHolder): Holder for original batch normalization statistics.
-        **kwargs: Additional arguments if needed.
+        activation_extractor (PytorchActivationExtractor): 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.
@@ -35,19 +40,25 @@ def average_bn_layer_weighting_fn(orig_bn_stats_holder: OriginalBNStatsHolder, *
     num_bn_layers = orig_bn_stats_holder.get_num_bn_layers()
     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: OriginalBNStatsHolder, **kwargs) -> Dict[str, float]:
+def first_bn_multiplier_weighting_fn(orig_bn_stats_holder: OriginalBNStatsHolder,
+                                  activation_extractor: PytorchActivationExtractor,
+                                  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 (OriginalBNStatsHolder): Holder for original batch normalization statistics.
-        **kwargs: Additional arguments if needed.
+        activation_extractor (PytorchActivationExtractor): 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.
     """
     layer_weighting_dict = {orig_bn_stats_holder.get_bn_layer_names()[0]: 10}
-    return layer_weighting_dict.update({bn_layer_name: 1  for bn_layer_name in orig_bn_stats_holder.get_bn_layer_names()[1:]})
+    layer_weighting_dict.update({bn_layer_name: 1 for bn_layer_name in orig_bn_stats_holder.get_bn_layer_names()[1:]})
+    return layer_weighting_dict
 
 
 # Dictionary of layer weighting functions

model_compression_toolkit/data_generation/pytorch/optimization_functions/image_initilization.py

@@ -104,7 +104,7 @@ def diverse_sample(size: Tuple[int, ...]) -> Tensor:
 
 def default_data_init_fn(
         n_images: int = 1000,
-        size: tuple = (224, 224),
+        size: Union[int, Tuple[int, int]] = (224, 224),
         crop: int = 32,
         sample_fn: Callable = diverse_sample,
         batch_size: int = 50) -> Tuple[int, DataLoader]:
@@ -113,7 +113,7 @@ def default_data_init_fn(
 
     Args:
         n_images (int): The number of random samples.
-        size (Tuple[int, int]): The size of each sample.
+        size (Union[int, Tuple[int, int]]): The size of each sample.
         crop (int): The crop size.
         sample_fn (Callable): The function to generate a random sample.
         batch_size (int): The batch size.

model_compression_toolkit/data_generation/pytorch/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 torch.optim.optimizer import Optimizer
+from torch import inf
+from typing import Union, List, Dict, Any
+
+from model_compression_toolkit.logger import Logger
+
+
+class ReduceLROnPlateauWithReset:
+    """
+    Reduce learning rate when a metric has stopped improving. This scheduler allows resetting
+    the learning rate to the initial value after a specified number of bad epochs.
+    """
+
+    def __init__(self, optimizer: Optimizer, mode: str = 'min', factor: float = 0.1, patience: int = 10,
+                 threshold: float = 1e-4, threshold_mode: str = 'rel', cooldown: int = 0,
+                 min_lr: Union[float, List[float]] = 0, eps: float = 1e-8, verbose: bool = False):
+        """
+        Initialize the ReduceLROnPlateauWithReset scheduler.
+
+        Args:
+            optimizer (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 or list): A scalar or a list of scalars. A lower bound on the learning rate of all param groups
+                                    or each group respectively. 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.
+        """
+        if factor >= 1.0:
+            Logger.critical('Factor should be < 1.0.') # pragma: no cover
+        self.factor = factor
+
+        # Attach optimizer
+        if not isinstance(optimizer, Optimizer):
+            Logger.critical('{} is not an Optimizer'.format(
+                type(optimizer).__name__)) # pragma: no cover
+        self.optimizer = optimizer
+
+        if isinstance(min_lr, (list, tuple)):
+            if len(min_lr) != len(optimizer.param_groups):
+                Logger.critical("expected {} min_lrs, got {}".format(
+                    len(optimizer.param_groups), len(min_lr))) # pragma: no cover
+            self.min_lrs = list(min_lr)
+        else:
+            self.min_lrs = [min_lr] * len(optimizer.param_groups)
+
+        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()
+        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 step(self, metrics: float, epoch: Union[int, None] = None) -> None:
+        """
+        Update learning rate based on the given metrics.
+
+        Args:
+            metrics (float): The value of the metric to evaluate.
+            epoch (int, optional): The current epoch number. If not provided, it is incremented.
+        """
+        # Convert `metrics` to float, in case it's a zero-dim Tensor
+        current = float(metrics)
+        if epoch is None:
+            epoch = self.last_epoch + 1
+        self.last_epoch = epoch
+
+        # Check if the current metrics are better than the best
+        if self.is_better(current, self.best):
+            self.best = current
+            self.num_bad_epochs = 0
+        else:
+            self.num_bad_epochs += 1
+
+        # Handle cooldown period
+        if self.in_cooldown:
+            self.cooldown_counter -= 1
+            self.num_bad_epochs = 0  # Ignore any bad epochs in cooldown
+
+        # Reduce learning rate if the number of bad epochs exceeds patience
+        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
+
+        self._last_lr = [group['lr'] for group in self.optimizer.param_groups]
+
+    def _reduce_lr(self, epoch: int) -> None:
+        """
+        Reduce the learning rate for each parameter group.
+
+        Args:
+            epoch (int): The current epoch number.
+        """
+        for i, param_group in enumerate(self.optimizer.param_groups):
+            old_lr = float(param_group['lr'])
+            new_lr = max(old_lr * self.factor, self.min_lrs[i])
+            if old_lr - new_lr > self.eps:
+                param_group['lr'] = new_lr
+                if self.verbose:
+                    epoch_str = ("%.2f" if isinstance(epoch, float) else "%.5d") % epoch
+                    print('Epoch {}: reducing learning rate'
+                          ' of group {} to {:.4e}.'.format(epoch_str, i, new_lr))
+
+    @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) -> None:
+        """
+        Initialize the comparison function for determining if a new value is better.
+
+        Raises:
+            ValueError: If an unknown mode or threshold mode is provided.
+        """
+        if self.mode not in {'min', 'max'}:
+            Logger.critical('mode ' + self.mode + ' is unknown!') # pragma: no cover
+        if self.threshold_mode not in {'rel', 'abs'}:
+            Logger.critical('threshold mode ' + self.threshold_mode + ' is unknown!') # pragma: no cover
+
+        if self.mode == 'min':
+            self.mode_worse = float('inf')
+        else:  # mode == 'max':
+            self.mode_worse = float('-inf')
+
+    def state_dict(self) -> Dict[str, Any]:
+        """
+        Return the state of the scheduler as a dictionary.
+
+        Returns:
+            dict: The state of the scheduler.
+        """
+        return {key: value for key, value in self.__dict__.items() if key != 'optimizer'}
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        """
+        Load the scheduler state.
+
+        Args:
+            state_dict (dict): The state dictionary to load.
+        """
+        self.__dict__.update(state_dict)
+        self._init_is_better()