quadra 0.0.1__py3-none-any.whl → 2.1.13__py3-none-any.whl

quadra/modules/ssl/common.py

@@ -0,0 +1,285 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from quadra.utils.models import trunc_normal_
+
+
+class ProjectionHead(torch.nn.Module):
+    """Base class for all projection and prediction heads.
+
+    Args:
+        blocks:
+            List of tuples, each denoting one block of the projection head MLP.
+            Each tuple reads (linear_layer, batch_norm_layer, non_linearity_layer).
+            `batch_norm` layer can be possibly None, the same happens for
+            `non_linearity_layer`.
+    """
+
+    def __init__(self, blocks: list[tuple[torch.nn.Module | None, ...]]):
+        super().__init__()
+
+        layers: list[nn.Module] = []
+        for linear, batch_norm, non_linearity in blocks:
+            if linear:
+                layers.append(linear)
+            if batch_norm:
+                layers.append(batch_norm)
+            if non_linearity:
+                layers.append(non_linearity)
+        self.layers = torch.nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor):
+        return self.layers(x)
+
+
+class ExpanderReducer(ProjectionHead):
+    """Expander followed by a reducer."""
+
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__(
+            [
+                (
+                    torch.nn.Linear(input_dim, hidden_dim, bias=False),
+                    torch.nn.BatchNorm1d(hidden_dim),
+                    torch.nn.ReLU(inplace=True),
+                ),
+                (
+                    torch.nn.Linear(hidden_dim, output_dim, bias=False),
+                    torch.nn.BatchNorm1d(output_dim, affine=False),
+                    torch.nn.ReLU(inplace=True),
+                ),
+            ]
+        )
+
+
+class BarlowTwinsProjectionHead(ProjectionHead):
+    """Projection head used for Barlow Twins.
+    "The projector network has three linear layers, each with 8192 output
+    units. The first two layers of the projector are followed by a batch
+    normalization layer and rectified linear units." https://arxiv.org/abs/2103.03230.
+    """
+
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__(
+            [
+                (
+                    torch.nn.Linear(input_dim, hidden_dim, bias=False),
+                    torch.nn.BatchNorm1d(hidden_dim),
+                    torch.nn.ReLU(inplace=True),
+                ),
+                (
+                    torch.nn.Linear(hidden_dim, hidden_dim, bias=False),
+                    torch.nn.BatchNorm1d(hidden_dim),
+                    torch.nn.ReLU(inplace=True),
+                ),
+                (torch.nn.Linear(hidden_dim, output_dim, bias=False), None, None),
+            ]
+        )
+
+
+class SimCLRProjectionHead(ProjectionHead):
+    """Projection head used for SimCLR.
+    "We use a MLP with one hidden layer to obtain zi = g(h) = W_2 * σ(W_1 * h)
+    where σ is a ReLU non-linearity." https://arxiv.org/abs/2002.05709.
+    """
+
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__(
+            [
+                (
+                    torch.nn.Linear(input_dim, hidden_dim),
+                    None,
+                    torch.nn.ReLU(inplace=True),
+                ),
+                (torch.nn.Linear(hidden_dim, output_dim), None, None),
+            ]
+        )
+
+
+class SimCLRPredictionHead(ProjectionHead):
+    """Prediction head used for SimCLR.
+    "We set g(h) = W(2)σ(W(1)h), with the same input and output dimensionality (i.e. 2048)."
+    https://arxiv.org/abs/2002.05709.
+    """
+
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__(
+            [
+                (
+                    torch.nn.Linear(input_dim, hidden_dim, bias=False),
+                    torch.nn.BatchNorm1d(hidden_dim),
+                    torch.nn.ReLU(inplace=True),
+                ),
+                (torch.nn.Linear(hidden_dim, output_dim, bias=False), None, None),
+            ]
+        )
+
+
+class SimSiamProjectionHead(ProjectionHead):
+    """Projection head used for SimSiam.
+    "The projection MLP (in f) has BN applied to each fully-connected (fc)
+    layer, including its output fc. Its output fc has no ReLU. The hidden fc is
+    2048-d. This MLP has 3 layers." https://arxiv.org/abs/2011.10566.
+    """
+
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__(
+            [
+                (
+                    torch.nn.Linear(input_dim, hidden_dim, bias=False),
+                    torch.nn.BatchNorm1d(hidden_dim),
+                    torch.nn.ReLU(inplace=True),
+                ),
+                (
+                    torch.nn.Linear(hidden_dim, hidden_dim, bias=False),
+                    torch.nn.BatchNorm1d(hidden_dim, affine=False),
+                    torch.nn.ReLU(inplace=True),
+                ),
+                (
+                    torch.nn.Linear(hidden_dim, output_dim, bias=False),
+                    torch.nn.BatchNorm1d(output_dim, affine=False),
+                    None,
+                ),
+            ]
+        )
+
+
+class SimSiamPredictionHead(ProjectionHead):
+    """Prediction head used for SimSiam.
+    "The prediction MLP (h) has BN applied to its hidden fc layers. Its output
+    fc does not have BN (...) or ReLU. This MLP has 2 layers." https://arxiv.org/abs/2011.10566.
+    """
+
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__(
+            [
+                (
+                    torch.nn.Linear(input_dim, hidden_dim, bias=False),
+                    torch.nn.BatchNorm1d(hidden_dim),
+                    torch.nn.ReLU(inplace=True),
+                ),
+                (torch.nn.Linear(hidden_dim, output_dim, bias=False), None, None),
+            ]
+        )
+
+
+class BYOLPredictionHead(ProjectionHead):
+    """Prediction head used for BYOL."""
+
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__(
+            [
+                (
+                    torch.nn.Linear(input_dim, hidden_dim, bias=False),
+                    torch.nn.BatchNorm1d(hidden_dim),
+                    torch.nn.ReLU(inplace=True),
+                ),
+                (torch.nn.Linear(hidden_dim, output_dim, bias=False), None, None),
+            ]
+        )
+
+
+class BYOLProjectionHead(ProjectionHead):
+    """Projection head used for BYOL."""
+
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__(
+            [
+                (
+                    torch.nn.Linear(input_dim, hidden_dim, bias=False),
+                    torch.nn.BatchNorm1d(hidden_dim),
+                    torch.nn.ReLU(inplace=True),
+                ),
+                (torch.nn.Linear(hidden_dim, output_dim, bias=False), None, None),
+            ]
+        )
+
+
+class DinoProjectionHead(nn.Module):
+    """Projection head used for Dino. This projection head does not have
+    a batch norm layer.
+
+    Args:
+        input_dim: Input dimension for MLP head.
+        output_dim: Output dimension (projection dimension) for MLP head.
+        hidden_dim: Hidden dimension. Defaults to 512.
+        bottleneck_dim: Bottleneck dimension. Defaults to 256.
+        num_layers: Number of hidden layers used in MLP. Defaults to 3.
+        norm_last_layer: Decides applying normalization before last layer.
+            Defaults to False.
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        hidden_dim: int,
+        use_bn: bool = False,
+        norm_last_layer: bool = True,
+        num_layers: int = 3,
+        bottleneck_dim: int = 256,
+    ):
+        super().__init__()
+        num_layers = max(num_layers, 1)
+        self.mlp: nn.Linear | nn.Sequential
+        if num_layers == 1:
+            self.mlp = nn.Linear(input_dim, bottleneck_dim)
+        else:
+            layers: list[nn.Module] = [nn.Linear(input_dim, hidden_dim)]
+            if use_bn:
+                layers.append(nn.BatchNorm1d(hidden_dim))
+            layers.append(nn.GELU())
+            for _ in range(num_layers - 2):
+                layers.append(nn.Linear(hidden_dim, hidden_dim))
+                if use_bn:
+                    layers.append(nn.BatchNorm1d(hidden_dim))
+                layers.append(nn.GELU())
+            layers.append(nn.Linear(hidden_dim, bottleneck_dim))
+            self.mlp = nn.Sequential(*layers)
+        self.apply(self._init_weights)
+        self.last_layer = nn.utils.weight_norm(nn.Linear(bottleneck_dim, output_dim, bias=False))
+        self.last_layer.weight_g.data.fill_(1)
+        if norm_last_layer:
+            self.last_layer.weight_g.requires_grad = False
+
+    def _init_weights(self, m):
+        """Initialize the weights of the projection head."""
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x = self.mlp(x)
+        x = nn.functional.normalize(x, dim=-1, p=2)
+        x = self.last_layer(x)
+        return x
+
+
+class MultiCropModel(nn.Module):
+    """MultiCrop model for DINO augmentation.
+
+    It takes 2 global crops and N (possible) local crops as a single tensor.
+
+    Args:
+        backbone: Backbone model.
+        head: Head model.
+    """
+
+    def __init__(self, backbone: nn.Module, head: nn.Module):
+        super().__init__()
+        self.backbone = backbone
+        self.head = head
+
+    def forward(self, x):
+        n_crops = len(x)
+        # (n_samples * n_crops, 3, size, size)
+        concatenated = torch.cat(x, dim=0)
+        # (n_samples * n_crops, in_dim)
+        embedding = self.backbone(concatenated)
+        logits = self.head(embedding)  # (n_samples * n_crops, out_dim)
+        chunks = logits.chunk(n_crops)  # n_crops * (n_samples, out_dim)
+
+        return chunks

quadra/modules/ssl/dino.py

@@ -0,0 +1,186 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import sklearn
+import torch
+from pytorch_lightning.core.optimizer import LightningOptimizer
+from torch import nn
+from torch.optim import Optimizer
+
+from quadra.modules.ssl import BYOL
+from quadra.utils.models import clip_gradients
+from quadra.utils.utils import get_logger
+
+log = get_logger(__name__)
+
+
+class Dino(BYOL):
+    """DINO pytorch-lightning module.
+
+    Args:
+        student : student model
+        teacher : teacher model
+        student_projection_mlp : student projection MLP
+        teacher_projection_mlp : teacher projection MLP
+        criterion : loss function
+        freeze_last_layer : number of layers to freeze in the student model. Default: 1
+        classifier: Standard sklearn classifier
+        optimizer: optimizer of the training. If None a default Adam is used.
+        lr_scheduler: lr scheduler. If None a default ReduceLROnPlateau is used.
+        lr_scheduler_interval: interval at which the lr scheduler is updated.
+        teacher_momentum: momentum of the teacher parameters
+        teacher_momentum_cosine_decay: whether to use cosine decay for the teacher momentum
+    """
+
+    def __init__(
+        self,
+        student: nn.Module,
+        teacher: nn.Module,
+        student_projection_mlp: nn.Module,
+        teacher_projection_mlp: nn.Module,
+        criterion: nn.Module,
+        freeze_last_layer: int = 1,
+        classifier: sklearn.base.ClassifierMixin | None = None,
+        optimizer: Optimizer | None = None,
+        lr_scheduler: object | None = None,
+        lr_scheduler_interval: str | None = "epoch",
+        teacher_momentum: float = 0.9995,
+        teacher_momentum_cosine_decay: bool | None = True,
+    ):
+        super().__init__(
+            student=student,
+            teacher=teacher,
+            student_projection_mlp=student_projection_mlp,
+            student_prediction_mlp=nn.Identity(),
+            teacher_projection_mlp=teacher_projection_mlp,
+            criterion=criterion,
+            teacher_momentum=teacher_momentum,
+            teacher_momentum_cosine_decay=teacher_momentum_cosine_decay,
+            classifier=classifier,
+            optimizer=optimizer,
+            lr_scheduler=lr_scheduler,
+            lr_scheduler_interval=lr_scheduler_interval,
+        )
+        self.freeze_last_layer = freeze_last_layer
+
+    def initialize_teacher(self):
+        """Initialize teacher from the state dict of the student one,
+        checking also that student model requires greadient correctly.
+        """
+        self.teacher_projection_mlp.load_state_dict(self.student_projection_mlp.state_dict())
+        for p in self.teacher_projection_mlp.parameters():
+            p.requires_grad = False
+
+        self.teacher.load_state_dict(self.model.state_dict())
+        for p in self.teacher.parameters():
+            p.requires_grad = False
+
+        all_frozen = True
+        for p in self.model.parameters():
+            all_frozen = all_frozen and (not p.requires_grad)
+
+        if all_frozen:
+            log.warning(
+                "All parameters of the student model are frozen, the model will not be trained, automatically"
+                " unfreezing all the layers"
+            )
+
+            for p in self.model.parameters():
+                p.requires_grad = True
+
+        for name, p in self.student_projection_mlp.named_parameters():
+            if name != "last_layer.weight_g":
+                assert p.requires_grad is True
+
+        self.teacher_initialized = True
+
+    def student_multicrop_forward(self, x: list[torch.Tensor]) -> torch.Tensor:
+        """Student forward on the multicrop imges.
+
+        Args:
+            x: List of torch.Tensor containing multicropped augmented images
+
+        Returns:
+            torch.Tensor: a tensor of shape NxBxD, where N is the number crops
+                corresponding to the length of the input list `x`, B is the batch size
+                and D is the output dimension
+        """
+        n_crops = len(x)
+        concatenated = torch.cat(x, dim=0)  # (n_samples * n_crops, C, H, W)
+        embedding = self.model(concatenated)  # (n_samples * n_crops, in_dim)
+        logits = self.student_projection_mlp(embedding)  # (n_samples * n_crops, out_dim)
+        chunks = logits.chunk(n_crops)  # n_crops * (n_samples, out_dim)
+        return chunks
+
+    def teacher_multicrop_forward(self, x: list[torch.Tensor]) -> torch.Tensor:
+        """Teacher forward on the multicrop imges.
+
+        Args:
+            x: List of torch.Tensor containing multicropped augmented images
+
+        Returns:
+            torch.Tensor: a tensor of shape NxBxD, where N is the number crops
+                corresponding to the length of the input list `x`, B is the batch size
+                and D is the output dimension
+        """
+        n_crops = len(x)
+        concatenated = torch.cat(x, dim=0)  # (n_samples * n_crops, C, H, W)
+        embedding = self.teacher(concatenated)  # (n_samples * n_crops, in_dim)
+        logits = self.teacher_projection_mlp(embedding)  # (n_samples * n_crops, out_dim)
+        chunks = logits.chunk(n_crops)  # n_crops * (n_samples, out_dim)
+        return chunks
+
+    def cancel_gradients_last_layer(self, epoch: int, freeze_last_layer: int):
+        """Zero out the gradient of the last layer, as specified in the paper.
+
+        Args:
+            epoch: current epoch
+            freeze_last_layer: maximum freeze epoch: if `epoch` >= `freeze_last_layer`
+                then the gradient of the last layer will not be freezed
+        """
+        if epoch >= freeze_last_layer:
+            return
+        for n, p in self.student_projection_mlp.named_parameters():
+            if "last_layer" in n:
+                p.grad = None
+
+    def training_step(self, batch: tuple[list[torch.Tensor], torch.Tensor], *args: Any) -> torch.Tensor:
+        images, _ = batch
+        with torch.no_grad():
+            teacher_output = self.teacher_multicrop_forward(images[:2])
+
+        student_output = self.student_multicrop_forward(images)
+        loss = self.criterion(self.current_epoch, student_output, teacher_output)
+
+        self.log(name="loss", value=loss, on_step=True, on_epoch=True, prog_bar=True)
+        return loss
+
+    def configure_gradient_clipping(
+        self,
+        optimizer: Optimizer,
+        gradient_clip_val: int | float | None = None,
+        gradient_clip_algorithm: str | None = None,
+    ):
+        """Configure gradient clipping for the optimizer."""
+        if gradient_clip_algorithm is not None and gradient_clip_val is not None:
+            clip_gradients(self.model, gradient_clip_val)
+            clip_gradients(self.student_projection_mlp, gradient_clip_val)
+        self.cancel_gradients_last_layer(self.current_epoch, self.freeze_last_layer)
+
+    def optimizer_step(
+        self,
+        epoch: int,
+        batch_idx: int,
+        optimizer: Optimizer | LightningOptimizer,
+        optimizer_closure: Callable[[], Any] | None = None,
+    ) -> None:
+        """Override optimizer step to update the teacher parameters."""
+        super().optimizer_step(
+            epoch,
+            batch_idx,
+            optimizer,
+            optimizer_closure=optimizer_closure,
+        )
+        self.update_teacher()

quadra/modules/ssl/hyperspherical.py

@@ -0,0 +1,206 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from enum import Enum
+
+import torch
+import torch.nn.functional as F
+from torch import nn, optim
+
+from quadra.losses.ssl import hyperspherical as loss
+from quadra.modules.base import BaseLightningModule
+
+
+class AlignLoss(Enum):
+    """Align loss enum."""
+
+    L2 = 1
+    COSINE = 2
+
+
+class TLHyperspherical(BaseLightningModule):
+    """Hyperspherical model: maps features extracted from a pretrained backbone into
+    an hypersphere.
+
+    Args:
+        model: Feature extractor as pytorch `torch.nn.Module`
+        optimizer: optimizer of the training.
+            If None a default Adam is used.
+        lr_scheduler: lr scheduler.
+            If None a default ReduceLROnPlateau is used.
+        align_weight: Weight for the align loss component for the
+            hyperspherical loss.
+            Defaults to 1.
+        unifo_weight: Weight for the uniform loss component for the
+            hyperspherical loss.
+            Defaults to 1.
+        classifier_weight: Weight for the classifier loss component for the
+            hyperspherical loss.
+            Defaults to 1.
+        align_loss_type: Which type of align loss to use.
+            Defaults to AlignLoss.L2.
+        classifier_loss: Whether to compute a classifier loss to 'enhance'
+            the hyperpsherical loss with the classification loss.
+            It True, model.classifier must be defined
+            Defaults to False.
+        num_classes: Number of classes for a classification problem.
+            Defaults to None.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        optimizer: optim.Optimizer | None = None,
+        lr_scheduler: object | None = None,
+        align_weight: float = 1,
+        unifo_weight: float = 1,
+        classifier_weight: float = 1,
+        align_loss_type: AlignLoss = AlignLoss.L2,
+        classifier_loss: bool = False,
+        num_classes: int | None = None,
+    ):
+        super().__init__(model, optimizer, lr_scheduler)
+        self.align_loss_fun: (
+            Callable[[torch.Tensor, torch.Tensor, int], torch.Tensor]
+            | Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
+        )
+        self.align_weight = align_weight
+        self.unifo_weight = unifo_weight
+        self.classifier_weight = classifier_weight
+        self.align_loss_type = align_loss_type
+        if align_loss_type == AlignLoss.L2:
+            self.align_loss_fun = loss.align_loss
+        elif align_loss_type == AlignLoss.COSINE:
+            self.align_loss_fun = loss.cosine_align_loss
+        else:
+            raise ValueError("The align loss must be one of 'AlignLoss.L2' (L2 distance) or AlignLoss.COSINE")
+
+        if classifier_loss and model.classifier is None:
+            raise AssertionError("Classifier is not defined")
+
+        self.classifier_loss = classifier_loss
+        self.num_classes = num_classes
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        # pylint: disable=unused-argument
+        im_x, im_y, target = batch
+        emb_x, emb_y = self(torch.cat([im_x, im_y])).chunk(2)
+
+        align_loss = 0.0
+        if self.align_weight > 0:
+            align_loss = self.align_loss_fun(emb_x, emb_y)
+
+        unifo_loss = 0.0
+        if self.unifo_weight > 0:
+            unifo_loss = (loss.uniform_loss(emb_x) + loss.uniform_loss(emb_y)) / 2
+
+        classifier_loss = 0.0
+        if self.classifier_loss:
+            pred = self.model.classifier(emb_x)
+            classifier_loss = F.cross_entropy(pred, target)
+
+        total_loss = (
+            self.align_weight * align_loss + self.unifo_weight * unifo_loss + self.classifier_weight * classifier_loss
+        )
+
+        self.log(
+            "t_loss",
+            total_loss,
+            on_epoch=True,
+            logger=True,
+            prog_bar=False,
+        )
+        self.log(
+            "t_align",
+            align_loss,
+            on_epoch=True,
+            on_step=False,
+            logger=True,
+            prog_bar=False,
+        )
+        self.log(
+            "t_classifier",
+            classifier_loss,
+            on_epoch=True,
+            on_step=False,
+            logger=True,
+            prog_bar=True,
+        )
+        self.log(
+            "t_unif",
+            unifo_loss,
+            on_epoch=True,
+            on_step=False,
+            logger=True,
+            prog_bar=False,
+        )
+        return {"loss": total_loss}
+
+    def train_epoch_end(self, outputs):
+        avg_loss = torch.stack([x["loss"] for x in outputs]).mean()
+
+        return {"loss": avg_loss}
+
+    def validation_step(self, batch, batch_idx):
+        # pylint: disable=unused-argument
+        im_x, im_y, target = batch
+        emb_x, emb_y = self(torch.cat([im_x, im_y])).chunk(2)
+
+        align_loss = 0.0
+        if self.align_weight > 0:
+            align_loss = self.align_loss_fun(emb_x, emb_y)
+
+        unifo_loss = 0.0
+        if self.unifo_weight > 0:
+            unifo_loss = (loss.uniform_loss(emb_x) + loss.uniform_loss(emb_y)) / 2
+
+        classifier_loss = 0.0
+        if self.classifier_loss:
+            pred = self.model.classifier(emb_x)
+            classifier_loss = F.cross_entropy(pred, target)
+
+        total_loss = (
+            self.align_weight * align_loss + self.unifo_weight * unifo_loss + self.classifier_weight * classifier_loss
+        )
+
+        self.log(
+            "val_loss",
+            total_loss,
+            on_epoch=True,
+            on_step=False,
+            logger=True,
+            prog_bar=False,
+        )
+        self.log(
+            "v_classifier",
+            classifier_loss,
+            on_epoch=True,
+            on_step=False,
+            logger=True,
+            prog_bar=True,
+        )
+        self.log(
+            "v_align",
+            align_loss,
+            on_epoch=True,
+            on_step=False,
+            logger=True,
+            prog_bar=False,
+        )
+        self.log(
+            "v_unif",
+            unifo_loss,
+            on_epoch=True,
+            on_step=False,
+            logger=True,
+            prog_bar=False,
+        )
+        return {"val_loss": total_loss}
+
+    def on_validation_epoch_end(self, outputs):
+        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
+
+        return {"val_loss": avg_loss}