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

quadra/modules/ssl/idmm.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import sklearn
+import timm
+import torch
+import torch.nn.functional as F
+
+from quadra.modules.base import SSLModule
+
+
+class IDMM(SSLModule):
+    """IDMM model.
+
+    Args:
+        model: backbone model
+        prediction_mlp: student prediction MLP
+        criterion: loss function
+        multiview_loss: whether to use the multiview loss as definied in https://arxiv.org/abs/2201.10728.
+            Defaults to True.
+        mixup_fn: the mixup/cutmix function to be applied to a batch of images.
+            Defaults to None.
+        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.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        prediction_mlp: torch.nn.Module,
+        criterion: torch.nn.Module,
+        multiview_loss: bool = True,
+        mixup_fn: timm.data.Mixup | None = None,
+        classifier: sklearn.base.ClassifierMixin | None = None,
+        optimizer: torch.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.prediction_mlp = prediction_mlp
+        self.mixup_fn = mixup_fn
+        self.multiview_loss = multiview_loss
+
+    def forward(self, x):
+        z = self.model(x)
+        p = self.prediction_mlp(z)
+        return z, p
+
+    def training_step(self, batch, batch_idx):
+        # pylint: disable=unused-argument
+        # Compute loss
+        if self.multiview_loss:
+            im_x, im_y, target = batch
+
+            # Contrastive loss
+            za, _ = self(im_x)
+            zb, _ = self(im_y)
+            za = F.normalize(za, dim=-1)
+            zb = F.normalize(zb, dim=-1)
+            s_aa = za.T @ za
+            s_ab = za.T @ zb
+            contrastive = (
+                torch.log(torch.exp(s_aa).sum(-1))
+                - torch.diagonal(s_aa)
+                + torch.log(torch.exp(s_ab).sum(-1))
+                - torch.diagonal(s_ab)
+            )
+
+            # Instance discrimination
+            if self.mixup_fn is not None:
+                im_x, target = self.mixup_fn(im_x, target)
+            _, pred = self(im_x)
+            loss = self.criterion(pred, target) + contrastive.mean()
+        else:
+            im_x, target = batch
+            if self.mixup_fn is not None:
+                im_x, target = self.mixup_fn(im_x, target)
+            pred = self(im_x)
+            loss = self.criterion(pred, target)
+
+        self.log(
+            "loss",
+            loss,
+            on_epoch=True,
+            on_step=True,
+            logger=True,
+            prog_bar=True,
+        )
+        return loss

quadra/modules/ssl/simclr.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import sklearn
+import torch
+from torch import nn
+
+from quadra.modules.base import SSLModule
+
+
+class SimCLR(SSLModule):
+    """SIMCLR class.
+
+    Args:
+        model: Feature extractor as pytorch `torch.nn.Module`
+        projection_mlp: projection head as
+            pytorch `torch.nn.Module`
+        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: torch.nn.Module,
+        classifier: sklearn.base.ClassifierMixin | None = None,
+        optimizer: torch.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.projection_mlp = projection_mlp
+
+    def forward(self, x):
+        x = self.model(x)
+        x = self.projection_mlp(x)
+        return x
+
+    def training_step(
+        self, batch: tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor], batch_idx: int
+    ) -> torch.Tensor:
+        """Args:
+            batch: The batch of data
+            batch_idx: The index of the batch.
+
+        Returns:
+            The computed loss
+        """
+        # pylint: disable=unused-argument
+        (im_x, im_y), _ = batch
+        emb_x = self(im_x)
+        emb_y = self(im_y)
+        loss = self.criterion(emb_x, emb_y)
+
+        self.log(
+            "loss",
+            loss,
+            on_epoch=True,
+            on_step=True,
+            logger=True,
+            prog_bar=True,
+        )
+        return loss

quadra/modules/ssl/simsiam.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import sklearn
+import torch
+
+from quadra.modules.base import SSLModule
+
+
+class SimSIAM(SSLModule):
+    """SimSIAM model.
+
+    Args:
+        model: Feature extractor as pytorch `torch.nn.Module`
+        projection_mlp: optional projection head as pytorch `torch.nn.Module`
+        prediction_mlp: optional predicition head as pytorch `torch.nn.Module`
+        criterion: loss to be applied.
+        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.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        projection_mlp: torch.nn.Module,
+        prediction_mlp: torch.nn.Module,
+        criterion: torch.nn.Module,
+        classifier: sklearn.base.ClassifierMixin | None = None,
+        optimizer: torch.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.prediction_mlp = prediction_mlp
+
+    def forward(self, x):
+        x = self.model(x)
+        z = self.projection_mlp(x)
+        p = self.prediction_mlp(z)
+        return p, z.detach()
+
+    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
+        p1, z1 = self(im_x)
+        p2, z2 = self(im_y)
+        loss = self.criterion(p1, p2, z1, z2)
+
+        self.log(
+            "loss",
+            loss,
+            on_epoch=True,
+            on_step=True,
+            logger=True,
+            prog_bar=True,
+        )
+        return loss

quadra/modules/ssl/vicreg.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import sklearn
+import torch
+from torch import nn, optim
+
+from quadra.modules.base import SSLModule
+
+
+class VICReg(SSLModule):
+    """VICReg 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/optimizers/__init__.py

@@ -0,0 +1,4 @@
+from .lars import LARS
+from .sam import SAM
+
+__all__ = ["LARS", "SAM"]

quadra/optimizers/lars.py

@@ -0,0 +1,153 @@
+"""References:
+- https://arxiv.org/pdf/1708.03888.pdf
+- https://github.com/pytorch/pytorch/blob/1.6/torch/optim/sgd.py.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+from torch.nn import Parameter
+from torch.optim.optimizer import Optimizer, _RequiredParameter, required
+
+
+class LARS(Optimizer):
+    r"""Extends SGD in PyTorch with LARS scaling from the paper
+    `Large batch training of Convolutional Networks <https://arxiv.org/pdf/1708.03888.pdf>`_.
+
+    Args:
+        params: iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr: learning rate
+        momentum: momentum factor (default: 0)
+        weight_decay: weight decay (L2 penalty) (default: 0)
+        dampening: dampening for momentum (default: 0)
+        nesterov: enables Nesterov momentum (default: False)
+        trust_coefficient: trust coefficient for computing LR (default: 0.001)
+        eps: eps for division denominator (default: 1e-8).
+
+    Example:
+        >>> model = torch.nn.Linear(10, 1)
+        >>> input = torch.Tensor(10)
+        >>> target = torch.Tensor([1.])
+        >>> loss_fn = lambda input, target: (input - target) ** 2
+        >>> #
+        >>> optimizer = LARS(model.parameters(), lr=0.1, momentum=0.9)
+        >>> optimizer.zero_grad()
+        >>> loss_fn(model(input), target).backward()
+        >>> optimizer.step()
+
+    .. note::
+        The application of momentum in the SGD part is modified according to
+        the PyTorch standards. LARS scaling fits into the equation in the
+        following fashion.
+
+        .. math::
+            \begin{aligned}
+                g_{t+1} & = \text{lars_lr} * (\beta * p_{t} + g_{t+1}), \\
+                v_{t+1} & = \\mu * v_{t} + g_{t+1}, \\
+                p_{t+1} & = p_{t} - \text{lr} * v_{t+1},
+            \\end{aligned}
+
+        where :math:`p`, :math:`g`, :math:`v`, :math:`\\mu` and :math:`\beta` denote the
+        parameters, gradient, velocity, momentum, and weight decay respectively.
+        The :math:`lars_lr` is defined by Eq. 6 in the paper.
+        The Nesterov version is analogously modified.
+
+    .. warning::
+        Parameters with weight decay set to 0 will automatically be excluded from
+        layer-wise LR scaling. This is to ensure consistency with papers like SimCLR
+        and BYOL.
+    """
+
+    def __init__(
+        self,
+        params: list[Parameter],
+        lr: _RequiredParameter = required,
+        momentum: float = 0,
+        dampening: float = 0,
+        weight_decay: float = 0,
+        nesterov: bool = False,
+        trust_coefficient: float = 0.001,
+        eps: float = 1e-8,
+    ):
+        if lr is not required and lr < 0.0:  # type: ignore[operator]
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if momentum < 0.0:
+            raise ValueError(f"Invalid momentum value: {momentum}")
+        if weight_decay < 0.0:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        defaults = {
+            "lr": lr,
+            "momentum": momentum,
+            "dampening": dampening,
+            "weight_decay": weight_decay,
+            "nesterov": nesterov,
+        }
+        if nesterov and (momentum <= 0 or dampening != 0):
+            raise ValueError("Nesterov momentum requires a momentum and zero dampening")
+
+        self.eps = eps
+        self.trust_coefficient = trust_coefficient
+
+        super().__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+
+        for group in self.param_groups:
+            group.setdefault("nesterov", False)
+
+    @torch.no_grad()
+    def step(self, closure: Callable | None = None):
+        """Performs a single optimization step.
+
+        Args:
+            closure: A closure that reevaluates the model and returns the loss. Defaults to None.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        # exclude scaling for params with 0 weight decay
+        for group in self.param_groups:
+            weight_decay = group["weight_decay"]
+            momentum = group["momentum"]
+            dampening = group["dampening"]
+            nesterov = group["nesterov"]
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                d_p = p.grad
+                p_norm = torch.norm(p.data)
+                g_norm = torch.norm(p.grad.data)
+
+                # lars scaling + weight decay part
+                if weight_decay != 0 and p_norm != 0 and g_norm != 0:
+                    lars_lr = p_norm / (g_norm + p_norm * weight_decay + self.eps)
+                    lars_lr *= self.trust_coefficient
+
+                    d_p = d_p.add(p, alpha=weight_decay)
+                    d_p *= lars_lr
+
+                # sgd part
+                if momentum != 0:
+                    param_state = self.state[p]
+                    if "momentum_buffer" not in param_state:
+                        buf = param_state["momentum_buffer"] = torch.clone(d_p).detach()
+                    else:
+                        buf = param_state["momentum_buffer"]
+                        buf.mul_(momentum).add_(d_p, alpha=1 - dampening)
+                    if nesterov:
+                        d_p = d_p.add(buf, alpha=momentum)
+                    else:
+                        d_p = buf
+
+                p.add_(d_p, alpha=-group["lr"])
+
+        return loss

quadra/optimizers/sam.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import torch
+from torch.nn import Parameter
+
+
+class SAM(torch.optim.Optimizer):
+    """PyTorch implementation of Sharpness-Aware-Minization paper: https://arxiv.org/abs/2010.01412
+    and https://arxiv.org/abs/2102.11600.
+    Taken from: https://github.com/davda54/sam.
+
+    Args:
+        params: model parameters.
+        base_optimizer: optimizer to use.
+        rho: Postive float value used to scale the gradients.
+        adaptive: Boolean flag indicating whether to use adaptive step update.
+        **kwargs: Additional parameters for the base optimizer.
+    """
+
+    def __init__(
+        self,
+        params: list[Parameter],
+        base_optimizer: torch.optim.Optimizer,
+        rho: float = 0.05,
+        adaptive: bool = True,
+        **kwargs: Any,
+    ):
+        assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"
+
+        defaults = {"rho": rho, "adaptive": adaptive, **kwargs}
+        super().__init__(params, defaults)
+
+        if callable(base_optimizer):
+            self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
+        else:
+            self.base_optimizer = base_optimizer
+        self.rho = rho
+        self.adaptive = adaptive
+        self.param_groups = self.base_optimizer.param_groups
+
+    @torch.no_grad()
+    def first_step(self, zero_grad: bool = False) -> None:
+        """First step for SAM optimizer.
+
+        Args:
+            zero_grad: Boolean flag indicating whether to zero the gradients.
+
+        Returns:
+            None
+        """
+        grad_norm = self._grad_norm()
+        for group in self.param_groups:
+            scale = self.rho / (grad_norm + 1e-12)
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                e_w = (torch.pow(p, 2) if self.adaptive else 1.0) * p.grad * scale.to(p)
+                p.add_(e_w)  # climb to the local maximum "w + e(w)"
+                self.state[p]["e_w"] = e_w
+
+        if zero_grad:
+            self.zero_grad()
+
+    @torch.no_grad()
+    def second_step(self, zero_grad: bool = False) -> None:
+        """Second step for SAM optimizer.
+
+        Args:
+            zero_grad: Boolean flag indicating whether to zero the gradients.
+
+        Returns:
+            None
+
+        """
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                p.sub_(self.state[p]["e_w"])  # get back to "w" from "w + e(w)"
+
+        self.base_optimizer.step()  # do the actual "sharpness-aware" update
+
+        if zero_grad:
+            self.zero_grad()
+
+    @torch.no_grad()
+    def step(self, closure: Callable | None = None) -> None:  # type: ignore[override]
+        """Step for SAM optimizer.
+
+        Args:
+            closure: The Optional closure for enable grad.
+
+        Returns:
+            None
+
+        """
+        if closure is not None:
+            closure = torch.enable_grad()(closure)
+
+        self.first_step(zero_grad=True)
+        if closure is not None:
+            closure()
+        self.second_step(zero_grad=False)
+
+    def _grad_norm(self) -> torch.Tensor:
+        """Put everything on the same device, in case of model parallelism
+        Returns:
+            Grad norm.
+        """
+        # put everything on the same device, in case of model parallelism
+        shared_device = self.param_groups[0]["params"][0].device
+        norm = torch.norm(
+            torch.stack(
+                [
+                    ((torch.abs(p) if self.adaptive else 1.0) * p.grad).norm(p=2).to(shared_device)
+                    for group in self.param_groups
+                    for p in group["params"]
+                    if p.grad is not None
+                ]
+            ),
+            p=2,
+        )
+        return norm

quadra/schedulers/__init__.py

@@ -0,0 +1,3 @@
+from .warmup import CosineAnnealingWithLinearWarmUp
+
+__all__ = ["CosineAnnealingWithLinearWarmUp"]

quadra/schedulers/base.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+class LearningRateScheduler(_LRScheduler):
+    """Provides inteface of learning rate scheduler.
+
+    Note:
+        Do not use this class directly, use one of the sub classes.
+    """
+
+    def __init__(self, optimizer: Optimizer, init_lr: tuple[float, ...]):
+        # pylint: disable=super-init-not-called
+        self.optimizer = optimizer
+        self.init_lr = init_lr
+
+    def step(self, *args, **kwargs):
+        """Base method, must be implemented by the sub classes."""
+        raise NotImplementedError
+
+    def set_lr(self, lr: tuple[float, ...]):
+        """Set the learning rate for the optimizer."""
+        if self.optimizer is not None:
+            for i, g in enumerate(self.optimizer.param_groups):
+                if "fix_lr" in g and g["fix_lr"]:
+                    if len(lr) == 1:
+                        lr_to_set = self.init_lr[0]
+                    else:
+                        lr_to_set = self.init_lr[i]
+                elif len(lr) == 1:
+                    lr_to_set = lr[0]
+                else:
+                    lr_to_set = lr[i]
+                g["lr"] = lr_to_set
+
+    def get_lr(self):
+        """Get the current learning rate if the optimizer is available."""
+        if self.optimizer is not None:
+            for g in self.optimizer.param_groups:
+                return g["lr"]
+
+        return None