Merge pull request #89 from discovery-unicamp/merging-from-dev-12052025

Merging changes from development version (12/05/2025)

Author: GabrielBG0

minerva/models/nets/mlp.py

@@ -1,73 +1,97 @@
-from torch import nn
-from typing import Sequence
+import torch.nn as nn
+from typing import Sequence, Optional, List
 
 
 class MLP(nn.Sequential):
     """
-    A multilayer perceptron (MLP) implemented as a subclass of nn.Sequential.
+    A flexible multilayer perceptron (MLP) implemented as a subclass of nn.Sequential.
 
-    This MLP is composed of a sequence of linear layers interleaved with ReLU activation
-    functions, except for the final layer which remains purely linear.
+    This class allows you to quickly build an MLP with:
+    - Custom layer sizes
+    - Configurable activation functions
+    - Optional intermediate operations (e.g., BatchNorm, Dropout) after each linear layer
+    - An optional final operation (e.g., normalization, final activation)
+
+    Parameters
+    ----------
+    layer_sizes : Sequence[int]
+        A list of integers specifying the sizes of each layer. Must contain at least two values:
+        the input and output dimensions.
+    activation_cls : type, optional
+        The activation function class (must inherit from nn.Module) to use between layers.
+        Defaults to nn.ReLU.
+    intermediate_ops : Optional[List[Optional[nn.Module]]], optional
+        A list of modules (e.g., nn.BatchNorm1d, nn.Dropout) to apply after each linear layer
+        and before the activation. Each item corresponds to one linear layer. Use `None` to skip
+        an operation for that layer. Must be the same length as the number of linear layers.
+    final_op : Optional[nn.Module], optional
+        A module to apply after the last layer (e.g., a final activation or normalization).
+
+    *args, **kwargs :
+        Additional arguments passed to the activation function constructor.
 
     Example
     -------
-
-    >>> mlp = MLP(10, 20, 30, 40)
+    >>> from torch import nn
+    >>> mlp = MLP(
+    ...     [128, 256, 64, 10],
+    ...     activation_cls=nn.ReLU,
+    ...     intermediate_ops=[nn.BatchNorm1d(256), nn.BatchNorm1d(64), None],
+    ...     final_op=nn.Sigmoid()
+    ... )
     >>> print(mlp)
     MLP(
-        (0): Linear(in_features=10, out_features=20, bias=True)
-        (1): ReLU()
-        (2): Linear(in_features=20, out_features=30, bias=True)
-        (3): ReLU()
-        (4): Linear(in_features=30, out_features=40, bias=True)
+        (0): Linear(in_features=128, out_features=256, bias=True)
+        (1): BatchNorm1d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+        (2): ReLU()
+        (3): Linear(in_features=256, out_features=64, bias=True)
+        (4): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+        (5): ReLU()
+        (6): Linear(in_features=64, out_features=10, bias=True)
+        (7): Sigmoid()
     )
     """
 
     def __init__(
         self,
         layer_sizes: Sequence[int],
         activation_cls: type = nn.ReLU,
+        intermediate_ops: Optional[List[Optional[nn.Module]]] = None,
+        final_op: Optional[nn.Module] = None,
         *args,
         **kwargs,
     ):
-        """
-        Initializes the MLP with specified layer sizes.
-
-        Parameters
-        ----------
-        layer_sizes : Sequence[int]
-            A sequence of positive integers indicating the size of each layer.
-            At least two integers are required, representing the input and output layers.
-        activation_cls : type
-            The class of the activation function to use between layers. Default is nn.ReLU.
-        *args
-            Additional arguments passed to the activation function.
-        **kwargs
-            Additional keyword arguments passed to the activation function.
-
-        Raises
-        ------
-        AssertionError
-            If fewer than two layer sizes are provided or if any layer size is not a positive integer.
-        AssertionError
-            If activation_cls does not inherit from torch.nn.Module.
-        """
 
         assert (
             len(layer_sizes) >= 2
         ), "Multilayer perceptron must have at least 2 layers"
         assert all(
-            ls > 0 and isinstance(ls, int) for ls in layer_sizes
+            isinstance(ls, int) and ls > 0 for ls in layer_sizes
         ), "All layer sizes must be positive integers"
-
         assert issubclass(
             activation_cls, nn.Module
         ), "activation_cls must inherit from torch.nn.Module"
 
+        num_layers = len(layer_sizes) - 1
+
+        if intermediate_ops is not None:
+            if len(intermediate_ops) != num_layers:
+                raise ValueError(
+                    f"Length of intermediate_ops ({len(intermediate_ops)}) must match number of layers ({num_layers})"
+                )
+
         layers = []
-        for i in range(len(layer_sizes) - 2):
-            layers.append(nn.Linear(layer_sizes[i], layer_sizes[i + 1]))
-            layers.append(activation_cls(*args, **kwargs))
-        layers.append(nn.Linear(layer_sizes[-2], layer_sizes[-1]))
+        for i in range(num_layers):
+            in_dim, out_dim = layer_sizes[i], layer_sizes[i + 1]
+            layers.append(nn.Linear(in_dim, out_dim))
+
+            if intermediate_ops is not None and intermediate_ops[i] is not None:
+                layers.append(intermediate_ops[i])
+
+            if activation_cls is not None:
+                layers.append(activation_cls(*args, **kwargs))
+
+        if final_op is not None:
+            layers.append(final_op)
 
         super().__init__(*layers)

minerva/models/ssl/byol.py

@@ -8,136 +8,82 @@
 from torch import nn
 from torch import Tensor
 from collections import OrderedDict
-from typing import List, Optional, Sequence, Tuple, Union
+from typing import Optional, Sequence
 
 from minerva.losses.negative_cossine_similatiry import NegativeCosineSimilarity
-
-# --- Model Parts ---------------------------------------------------------
-
-# Borrowed from https://github.com/lightly-ai/lightly/blob/master/lightly/models/modules/heads.py#L15
-
-
-class ProjectionHead(nn.Module):
-    """Base class for all projection and prediction heads."""
-
-    def __init__(
-        self,
-        blocks: Sequence[
-            Union[
-                Tuple[int, int, Optional[nn.Module], Optional[nn.Module]],
-                Tuple[int, int, Optional[nn.Module], Optional[nn.Module], bool],
-            ],
-        ],
-    ) -> None:
-        super().__init__()
-
-        layers: List[nn.Module] = []
-        for block in blocks:
-            input_dim, output_dim, batch_norm, non_linearity, *bias = block
-            use_bias = bias[0] if bias else not bool(batch_norm)
-            layers.append(nn.Linear(input_dim, output_dim, bias=use_bias))
-            if batch_norm:
-                layers.append(batch_norm)
-            if non_linearity:
-                layers.append(non_linearity)
-        self.layers = nn.Sequential(*layers)
-
-    def preprocess_step(self, x: Tensor) -> Tensor:
-        return x
-
-    def forward(self, x: Tensor) -> Tensor:
-        x = self.preprocess_step(x)
-        projection: Tensor = self.layers(x)
-        return projection
-
-
-class BYOLProjectionHead(ProjectionHead):
-    """Projection head used for BYOL.
-    "This MLP consists in a linear layer with output size 4096 followed by
-    batch normalization, rectified linear units (ReLU), and a final
-    linear layer with output dimension 256." [0]
-    [0]: BYOL, 2020, https://arxiv.org/abs/2006.07733
-    """
-
-    def __init__(
-        self, input_dim: int = 2048, hidden_dim: int = 4096, output_dim: int = 256
-    ):
-        super(BYOLProjectionHead, self).__init__(
-            [
-                (input_dim, hidden_dim, nn.BatchNorm1d(hidden_dim), nn.ReLU()),
-                (hidden_dim, output_dim, None, None),
-            ]
-        )
-
-        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
-
-    def preprocess_step(self, x: Tensor) -> Tensor:
-        return self.avgpool(x).flatten(start_dim=1)
-
-
-class BYOLPredictionHead(ProjectionHead):
-    """Prediction head used for BYOL.
-    "This MLP consists in a linear layer with output size 4096 followed by
-    batch normalization, rectified linear units (ReLU), and a final
-    linear layer with output dimension 256." [0]
-    [0]: BYOL, 2020, https://arxiv.org/abs/2006.07733
-    """
-
-    def __init__(
-        self, input_dim: int = 256, hidden_dim: int = 4096, output_dim: int = 256
-    ):
-        super(BYOLPredictionHead, self).__init__(
-            [
-                (input_dim, hidden_dim, nn.BatchNorm1d(hidden_dim), nn.ReLU()),
-                (hidden_dim, output_dim, None, None),
-            ]
-        )
-
-
-# --- Class implementation ----------------------------------------------------------
+from minerva.models.nets.mlp import MLP
+from torch.optim import Optimizer
+from minerva.models.nets.image.deeplabv3 import DeepLabV3Backbone
 
 
 class BYOL(L.LightningModule):
-    """A Bootstrap Your Own Latent (BYOL) model for self-supervised learning.
+    """Bootstrap Your Own Latent (BYOL) model for self-supervised learning.
 
     References
     ----------
     Grill, J., Strub, F., Altché, F., Tallec, C., Richemond, P. H., Buchatskaya, E., ... & Valko, M. (2020).
-    "Bootstrap your own latent-a new approach to self-supervised learning." Advances in neural information processing systems, 33, 21271-21284.
+    "Bootstrap your own latent - a new approach to self-supervised learning." Advances in Neural Information Processing Systems, 33, 21271-21284.
     """
 
     def __init__(
         self,
         backbone: Optional[nn.Module] = None,
-        learning_rate: float = 0.025,
-        schedule: int = 90000,
+        projection_head: Optional[nn.Module] = None,
+        prediction_head: Optional[nn.Module] = None,
+        learning_rate: Optional[float] = 1e-3,
+        schedule: Optional[int] = 90000,
+        criterion: Optional[Optimizer] = None,
     ):
         """
         Initializes the BYOL model.
 
         Parameters
         ----------
-        backbone: Optional[nn.Module]
-            The backbone network for feature extraction. Defaults to ResNet18.
-        learning_rate: float
-            The learning rate for the optimizer. Defaults to 0.025.
-        schedule: int
+        backbone : Optional[nn.Module]
+            The backbone network for feature extraction. Defaults to DeepLabV3Backbone.
+        projection_head : Optional[nn.Module]
+            Optional custom projection head module. If None, a default MLP-based projection head is used.
+        prediction_head : Optional[nn.Module]
+            Optional custom prediction head module. If None, a default MLP-based prediction head is used.
+        learning_rate : float
+            The learning rate for the optimizer. Defaults to 1e-3.
+        schedule : int
             The total number of steps for cosine decay scheduling. Defaults to 90000.
+        criterion : Optional[Optimizer]
+            Loss function to use. Defaults to NegativeCosineSimilarity.
         """
         super().__init__()
-        self.backbone = backbone or nn.Sequential(
-            *list(torchvision.models.resnet18().children())[:-1]
-        )
+        self.backbone = backbone or DeepLabV3Backbone()
         self.learning_rate = learning_rate
-        self.projection_head = BYOLProjectionHead(2048, 4096, 256)
-        self.prediction_head = BYOLPredictionHead(256, 4096, 256)
+        self.projection_head = projection_head or self._default_projection_head()
+        self.prediction_head = prediction_head or self._default_prediction_head()
         self.backbone_momentum = copy.deepcopy(self.backbone)
         self.projection_head_momentum = copy.deepcopy(self.projection_head)
         self.deactivate_requires_grad(self.backbone_momentum)
         self.deactivate_requires_grad(self.projection_head_momentum)
-        self.criterion = NegativeCosineSimilarity()
+        self.criterion = criterion or NegativeCosineSimilarity()
         self.schedule_length = schedule
 
+    def _default_projection_head(self) -> nn.Module:
+        """Creates the default projection head used in BYOL."""
+        return nn.Sequential(
+            nn.AdaptiveAvgPool2d((1, 1)),
+            nn.Flatten(start_dim=1),
+            MLP(
+                layer_sizes=[2048, 4096, 256],
+                activation_cls=nn.ReLU,
+                intermediate_ops=[nn.BatchNorm1d(4096), None],
+            ),
+        )
+
+    def _default_prediction_head(self) -> nn.Module:
+        """Creates the default prediction head used in BYOL."""
+        return MLP(
+            layer_sizes=[256, 4096, 256],
+            activation_cls=nn.ReLU,
+            intermediate_ops=[nn.BatchNorm1d(4096), None],
+        )
+
     def forward(self, x: Tensor) -> Tensor:
         """
         Forward pass for the BYOL model.

minerva/transforms/random_transform.py

@@ -170,15 +170,28 @@ def __init__(
         num_samples: int = 1,
         seed: Optional[int] = None,
     ):
+        """
+        Randomly applies a rotation to the image with a specified probability.
 
-        super().__init__(num_samples, seed)
-        self.prob = prob
+        Parameters
+        ----------
+        degrees : float
+            Maximum absolute value of the rotation angle in degrees. The angle is sampled
+            uniformly from [-degrees, +degrees].
+        prob : float
+            Probability that the rotation will be applied.
+        num_samples : int, optional
+            Number of samples to generate per call (for contrastive learning), default is 1.
+        seed : int, optional
+            Seed for the random number generator, useful for reproducibility.
+        """
+        super().__init__(num_samples=num_samples, seed=seed)
         self.degrees = degrees
+        self.prob = prob
 
     def select_transform(self):
-
         if self.rng.random() < self.prob:
-            degrees = self.rng.uniform(-self.degrees, self.degrees)
-            return Rotation(degrees=degrees)
+            angle = self.rng.uniform(-self.degrees, self.degrees)
+            return Rotation(degrees=angle)
         else:
             return EmptyTransform()