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

quadra/modules/classification/base.py

@@ -0,0 +1,327 @@
+from __future__ import annotations
+
+from typing import Any, cast
+
+import numpy as np
+import timm
+import torch
+import torchmetrics
+import torchmetrics.functional as TMF
+from pytorch_grad_cam import GradCAM
+from scipy import ndimage
+from torch import nn, optim
+
+from quadra.models.classification import BaseNetworkBuilder
+from quadra.modules.base import BaseLightningModule
+from quadra.utils.models import is_vision_transformer
+from quadra.utils.utils import get_logger
+from quadra.utils.vit_explainability import VitAttentionGradRollout
+
+log = get_logger(__name__)
+
+
+class ClassificationModule(BaseLightningModule):
+    """Lightning module for classification tasks.
+
+    Args:
+        model: Feature extractor as PyTorch `torch.nn.Module`
+        criterion: the loss to be applied as a PyTorch `torch.nn.Module`.
+        optimizer: optimizer of the training. Defaults to None.
+        lr_scheduler: Pytorch learning rate scheduler.
+            If None a default ReduceLROnPlateau is used.
+            Defaults to None.
+        lr_scheduler_interval: the learning rate scheduler interval.
+            Defaults to "epoch".
+        gradcam (bool): Whether to compute gradcam during prediction step
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        criterion: nn.Module,
+        optimizer: None | optim.Optimizer = None,
+        lr_scheduler: None | object = None,
+        lr_scheduler_interval: str | None = "epoch",
+        gradcam: bool = False,
+    ):
+        super().__init__(model, optimizer, lr_scheduler, lr_scheduler_interval)
+
+        self.criterion = criterion
+        self.gradcam = gradcam
+        self.train_acc = torchmetrics.Accuracy()
+        self.val_acc = torchmetrics.Accuracy()
+        self.test_acc = torchmetrics.Accuracy()
+        self.cam: GradCAM | None = None
+        self.grad_rollout: VitAttentionGradRollout | None = None
+
+        if not isinstance(self.model.features_extractor, timm.models.resnet.ResNet) and not is_vision_transformer(
+            cast(BaseNetworkBuilder, self.model).features_extractor
+        ):
+            log.warning(
+                "Backbone not compatible with gradcam. Only timm ResNets, timm ViTs and TorchHub dinoViTs supported",
+            )
+            self.gradcam = False
+
+        self.original_requires_grads: list[bool] = []
+
+    def forward(self, x: torch.Tensor):
+        return self.model(x)
+
+    def training_step(self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int):
+        # pylint: disable=unused-argument
+        im, target = batch
+        outputs = self(im)
+        loss = self.criterion(outputs, target)
+
+        self.log_dict(
+            {"train_loss": loss},
+            on_epoch=True,
+            on_step=True,
+            logger=True,
+            prog_bar=True,
+        )
+        self.log_dict(
+            {"train_acc": self.train_acc(outputs.argmax(1), target)},
+            on_step=False,
+            on_epoch=True,
+            logger=True,
+            prog_bar=True,
+        )
+        return loss
+
+    def validation_step(self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int):
+        # pylint: disable=unused-argument
+        im, target = batch
+        outputs = self(im)
+        loss = self.criterion(outputs, target)
+
+        self.log_dict(
+            {"val_loss": loss},
+            on_epoch=True,
+            on_step=True,
+            logger=True,
+            prog_bar=True,
+        )
+        self.log_dict(
+            {"val_acc": self.val_acc(outputs.argmax(1), target)},
+            on_step=False,
+            on_epoch=True,
+            logger=True,
+            prog_bar=True,
+        )
+        return loss
+
+    def test_step(self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int):
+        # pylint: disable=unused-argument
+        im, target = batch
+        outputs = self(im)
+
+        loss = self.criterion(outputs, target)
+
+        self.log_dict(
+            {"test_loss": loss},
+            on_epoch=True,
+            on_step=True,
+            logger=True,
+            prog_bar=False,
+        )
+        self.log_dict(
+            {"test_acc": self.test_acc(outputs.argmax(1), target)},
+            on_step=False,
+            on_epoch=True,
+            logger=True,
+            prog_bar=False,
+        )
+
+    def prepare_gradcam(self) -> None:
+        """Instantiate gradcam handlers."""
+        if isinstance(self.model.features_extractor, timm.models.resnet.ResNet):
+            target_layers = [cast(BaseNetworkBuilder, self.model).features_extractor.layer4[-1]]
+
+            self.cam = GradCAM(
+                model=self.model,
+                target_layers=target_layers,
+            )
+            # Activating gradients
+            for p in self.model.features_extractor.layer4[-1].parameters():
+                p.requires_grad = True
+        elif is_vision_transformer(cast(BaseNetworkBuilder, self.model).features_extractor):
+            self.grad_rollout = VitAttentionGradRollout(self.model)
+        else:
+            log.warning("Gradcam not implemented for this backbone, it won't be computed")
+            self.original_requires_grads.clear()
+            self.gradcam = False
+
+    def on_predict_start(self) -> None:
+        """If gradcam, prepares gradcam and saves params requires_grad state."""
+        if self.gradcam:
+            # Saving params requires_grad state
+            for p in self.model.parameters():
+                self.original_requires_grads.append(p.requires_grad)
+            self.prepare_gradcam()
+
+        return super().on_predict_start()
+
+    def on_predict_end(self) -> None:
+        """If we computed gradcam, requires_grad values are reset to original value."""
+        if self.gradcam:
+            # Get back to initial state
+            for i, p in enumerate(self.model.parameters()):
+                p.requires_grad = self.original_requires_grads[i]
+
+            # We are using GradCAM package only for resnets at the moment
+            if isinstance(self.model.features_extractor, timm.models.resnet.ResNet) and self.cam is not None:
+                # Needed to solve jitting bug
+                self.cam.activations_and_grads.release()
+            elif (
+                is_vision_transformer(cast(BaseNetworkBuilder, self.model).features_extractor)
+                and self.grad_rollout is not None
+            ):
+                for handle in self.grad_rollout.f_hook_handles:
+                    handle.remove()
+                for handle in self.grad_rollout.b_hook_handles:
+                    handle.remove()
+
+        return super().on_predict_end()
+
+    # pylint: disable=unused-argument
+    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
+        """Prediction step.
+
+        Args:
+            batch: Tuple composed by (image, target)
+            batch_idx: Batch index
+            dataloader_idx: Dataloader index
+        Returns:
+            Tuple containing:
+                predicted_classes: indexes of predicted classes
+                grayscale_cam: gray scale gradcams
+        """
+        im, _ = batch
+        outputs = self(im)
+        probs = torch.softmax(outputs, dim=1)
+        predicted_classes = torch.max(probs, dim=1).indices.tolist()
+        if self.gradcam:
+            # inference_mode set to false because gradcam needs gradients
+            with torch.inference_mode(False):
+                im = im.clone()
+
+                if isinstance(self.model.features_extractor, timm.models.resnet.ResNet) and self.cam:
+                    grayscale_cam = self.cam(input_tensor=im, targets=None)
+                elif (
+                    is_vision_transformer(cast(BaseNetworkBuilder, self.model).features_extractor) and self.grad_rollout
+                ):
+                    grayscale_cam_low_res = self.grad_rollout(input_tensor=im, targets_list=predicted_classes)
+                    orig_shape = grayscale_cam_low_res.shape
+                    new_shape = (orig_shape[0], im.shape[2], im.shape[3])
+                    zoom_factors = tuple(np.array(new_shape) / np.array(orig_shape))
+                    grayscale_cam = ndimage.zoom(grayscale_cam_low_res, zoom_factors, order=1)
+        else:
+            grayscale_cam = None
+        return predicted_classes, grayscale_cam, torch.max(probs, dim=1)[0].tolist()
+
+
+class MultilabelClassificationModule(BaseLightningModule):
+    """SklearnClassification model: train a generic SklearnClassification model for a multilabel
+    problem.
+
+    Args:
+        model: Feature extractor as PyTorch `torch.nn.Module`
+        criterion: the loss to be applied as a PyTorch `torch.nn.Module`.
+        optimizer: optimizer of the training. Defaults to None.
+        lr_scheduler: Pytorch learning rate scheduler.
+            If None a default ReduceLROnPlateau is used.
+            Defaults to None.
+        lr_scheduler_interval: the learning rate scheduler interval.
+            Defaults to "epoch".
+    """
+
+    def __init__(
+        self,
+        model: nn.Sequential,
+        criterion: nn.Module,
+        optimizer: None | optim.Optimizer = None,
+        lr_scheduler: None | object = None,
+        lr_scheduler_interval: str | None = "epoch",
+        gradcam: bool = False,
+    ):
+        super().__init__(model, optimizer, lr_scheduler, lr_scheduler_interval)
+        self.criterion = criterion
+        self.gradcam = gradcam
+
+        # TODO: can we use gradcam with more backbones?
+        if self.gradcam:
+            if not isinstance(model[0].features_extractor, timm.models.resnet.ResNet):
+                log.warning(
+                    "Backbone must be compatible with gradcam, at the moment only ResNets supported, disabling gradcam"
+                )
+                self.gradcam = False
+            else:
+                target_layers = [model[0].features_extractor.layer4[-1]]
+                self.cam = GradCAM(model=model, target_layers=target_layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        # pylint: disable=unused-argument
+        im, target = batch
+        outputs = self(im)
+        with torch.no_grad():
+            outputs_sig = torch.sigmoid(outputs)
+        loss = self.criterion(outputs, target)
+
+        self.log_dict(
+            {
+                "t_loss": loss,
+                "t_map": TMF.label_ranking_average_precision(outputs_sig, target.bool()),
+                "t_f1": TMF.f1_score(outputs_sig, target.bool(), average="samples"),
+            },
+            on_epoch=True,
+            on_step=False,
+            logger=True,
+            prog_bar=True,
+        )
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        # pylint: disable=unused-argument
+        im, target = batch
+        outputs = self(im)
+        with torch.no_grad():
+            outputs_sig = torch.sigmoid(outputs)
+        loss = self.criterion(outputs, target)
+
+        self.log_dict(
+            {
+                "val_loss": loss,
+                "val_map": TMF.label_ranking_average_precision(outputs_sig, target.bool()),
+                "val_f1": TMF.f1_score(outputs_sig, target.bool(), average="samples"),
+            },
+            on_epoch=True,
+            on_step=False,
+            logger=True,
+            prog_bar=True,
+        )
+        return loss
+
+    def test_step(self, batch, batch_idx):
+        # pylint: disable=unused-argument
+        im, target = batch
+        outputs = self(im)
+        with torch.no_grad():
+            outputs_sig = torch.sigmoid(outputs)
+        loss = self.criterion(outputs, target)
+
+        self.log_dict(
+            {
+                "test_loss": loss,
+                "test_map": TMF.label_ranking_average_precision(outputs_sig, target.bool()),
+                "test_f1": TMF.f1_score(outputs_sig, target.bool(), average="samples"),
+            },
+            on_epoch=True,
+            on_step=True,
+            logger=True,
+            prog_bar=False,
+        )
+        return loss

quadra/modules/ssl/__init__.py

@@ -0,0 +1,17 @@
+from .barlowtwins import BarlowTwins
+from .byol import BYOL
+from .dino import Dino
+from .idmm import IDMM
+from .simclr import SimCLR
+from .simsiam import SimSIAM
+from .vicreg import VICReg
+
+__all__ = [
+    "BarlowTwins",
+    "BYOL",
+    "Dino",
+    "IDMM",
+    "SimCLR",
+    "SimSIAM",
+    "VICReg",
+]

quadra/modules/ssl/barlowtwins.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import sklearn
+import torch
+from torch import nn, optim
+
+from quadra.modules.base import SSLModule
+
+
+class BarlowTwins(SSLModule):
+    """BarlowTwins model.
+
+    Args:
+        model: Network Module used for extract features
+        projection_mlp: Module to project extracted features
+        criterion: SSL loss to be applied
+        classifier: Standard sklearn classifier. Defaults to None.
+        optimizer: optimizer of the training. If None a default Adam is used. Defaults to None.
+        lr_scheduler: lr scheduler. If None a default ReduceLROnPlateau is used. Defaults to None.
+        lr_scheduler_interval: interval at which the lr scheduler is updated. Defaults to "epoch".
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        projection_mlp: nn.Module,
+        criterion: nn.Module,
+        classifier: sklearn.base.ClassifierMixin | None = None,
+        optimizer: optim.Optimizer | None = None,
+        lr_scheduler: object | None = None,
+        lr_scheduler_interval: str | None = "epoch",
+    ):
+        super().__init__(model, criterion, classifier, optimizer, lr_scheduler, lr_scheduler_interval)
+        # self.save_hyperparameters()
+        self.projection_mlp = projection_mlp
+        self.criterion = criterion
+
+    def forward(self, x):
+        x = self.model(x)
+        z = self.projection_mlp(x)
+        return z
+
+    def training_step(self, batch: tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor], batch_idx: int):
+        # pylint: disable=unused-argument
+        # Compute loss
+        (im_x, im_y), _ = batch
+        z1 = self(im_x)
+        z2 = self(im_y)
+        loss = self.criterion(z1, z2)
+
+        self.log(
+            "loss",
+            loss,
+            on_epoch=True,
+            on_step=True,
+            logger=True,
+            prog_bar=True,
+        )
+        return loss

quadra/modules/ssl/byol.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import math
+from collections.abc import Callable, Sized
+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.base import SSLModule
+
+
+class BYOL(SSLModule):
+    """BYOL module, inspired by https://arxiv.org/abs/2006.07733.
+
+    Args:
+        student : student model.
+        teacher : teacher model.
+        student_projection_mlp : student projection MLP.
+        student_prediction_mlp : student prediction MLP.
+        teacher_projection_mlp : teacher projection MLP.
+        criterion : loss function.
+        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. Default: True
+    """
+
+    def __init__(
+        self,
+        student: nn.Module,
+        teacher: nn.Module,
+        student_projection_mlp: nn.Module,
+        student_prediction_mlp: nn.Module,
+        teacher_projection_mlp: nn.Module,
+        criterion: nn.Module,
+        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__(
+            model=student,
+            criterion=criterion,
+            classifier=classifier,
+            optimizer=optimizer,
+            lr_scheduler=lr_scheduler,
+            lr_scheduler_interval=lr_scheduler_interval,
+        )
+        # Student model
+        self.max_steps: int
+        self.student_projection_mlp = student_projection_mlp
+        self.student_prediction_mlp = student_prediction_mlp
+
+        # Teacher model
+        self.teacher = teacher
+        self.teacher_projection_mlp = teacher_projection_mlp
+        self.teacher_initialized = False
+        self.teacher_momentum = teacher_momentum
+        self.teacher_momentum_cosine_decay = teacher_momentum_cosine_decay
+
+        self.initialize_teacher()
+
+    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
+
+        for p in self.student_projection_mlp.parameters():
+            assert p.requires_grad is True
+        for p in self.student_prediction_mlp.parameters():
+            assert p.requires_grad is True
+
+        self.teacher_initialized = True
+
+    def update_teacher(self):
+        """Update teacher given `self.teacher_momentum` by an exponential moving average
+        of the student parameters, that is: theta_t * tau + theta_s * (1 - tau), where
+        `theta_{s,t}` are the parameters of the student and the teacher model, while `tau` is the
+        teacher momentum. If `self.teacher_momentum_cosine_decay` is True, then the teacher
+        momentum will follow a cosine scheduling from `self.teacher_momentum` to 1:
+        tau = 1 - (1 - tau) * (cos(pi * t / T) + 1) / 2, where `t` is the current step and
+        `T` is the max number of steps.
+        """
+        with torch.no_grad():
+            if self.teacher_momentum_cosine_decay:
+                teacher_momentum = (
+                    1
+                    - (1 - self.teacher_momentum)
+                    * (math.cos(math.pi * self.trainer.global_step / self.max_steps) + 1)
+                    / 2
+                )
+            else:
+                teacher_momentum = self.teacher_momentum
+            self.log("teacher_momentum", teacher_momentum, prog_bar=True)
+            for student_ps, teacher_ps in zip(
+                list(self.model.parameters()) + list(self.student_projection_mlp.parameters()),
+                list(self.teacher.parameters()) + list(self.teacher_projection_mlp.parameters()),
+            ):
+                teacher_ps.data = teacher_ps.data * teacher_momentum + (1 - teacher_momentum) * student_ps.data
+
+    def on_train_start(self) -> None:
+        if isinstance(self.trainer.train_dataloader, Sized) and isinstance(self.trainer.max_epochs, int):
+            self.max_steps = len(self.trainer.train_dataloader) * self.trainer.max_epochs
+        else:
+            raise ValueError("BYOL requires `max_epochs` to be set and `train_dataloader` to be initialized.")
+
+    def training_step(self, batch: tuple[list[torch.Tensor], torch.Tensor], *args: Any) -> torch.Tensor:
+        [image1, image2], _ = batch
+
+        online_pred_one = self.student_prediction_mlp(self.student_projection_mlp(self.model(image1)))
+        online_pred_two = self.student_prediction_mlp(self.student_projection_mlp(self.model(image2)))
+
+        with torch.no_grad():
+            target_proj_one = self.teacher_projection_mlp(self.teacher(image1))
+            target_proj_two = self.teacher_projection_mlp(self.teacher(image2))
+
+        loss_one = self.criterion(online_pred_one, target_proj_two.detach())
+        loss_two = self.criterion(online_pred_two, target_proj_one.detach())
+        loss = loss_one + loss_two
+
+        self.log(name="loss", value=loss, on_step=True, on_epoch=True, prog_bar=True)
+        return loss
+
+    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()
+
+    def calculate_accuracy(self, batch):
+        """Calculate accuracy on the given batch."""
+        images, labels = batch
+        embedding = self.model(images).detach().cpu().numpy()
+        predictions = self.classifier.predict(embedding)
+        labels = labels.detach()
+        acc = self.val_acc(torch.tensor(predictions, device=self.device), labels)
+
+        return acc
+
+    def on_test_epoch_start(self) -> None:
+        self.fit_estimator()
+
+    def test_step(self, batch, *args: list[Any]) -> None:
+        """Calculate accuracy on the test set for the given batch."""
+        acc = self.calculate_accuracy(batch)
+        self.log(name="test_acc", value=acc, on_step=False, on_epoch=True, prog_bar=True)
+        return acc