careamics 0.0.19__py3-none-any.whl

careamics/losses/fcn/losses.py

@@ -0,0 +1,125 @@
+"""
+Loss submodule.
+
+This submodule contains the various losses used in CAREamics.
+"""
+
+import torch
+from torch.nn import L1Loss, MSELoss
+
+from careamics.models.lvae.noise_models import GaussianMixtureNoiseModel
+
+
+def mse_loss(source: torch.Tensor, target: torch.Tensor, *args) -> torch.Tensor:
+    """
+    Mean squared error loss.
+
+    Parameters
+    ----------
+    source : torch.Tensor
+        Source patches.
+    target : torch.Tensor
+        Target patches.
+    *args : Any
+        Additional arguments.
+
+    Returns
+    -------
+    torch.Tensor
+        Loss value.
+    """
+    loss = MSELoss()
+    return loss(source, target)
+
+
+def n2v_loss(
+    manipulated_batch: torch.Tensor,
+    original_batch: torch.Tensor,
+    masks: torch.Tensor,
+    *args,
+) -> torch.Tensor:
+    """
+    N2V Loss function described in A Krull et al 2018.
+
+    Parameters
+    ----------
+    manipulated_batch : torch.Tensor
+        Batch after manipulation function applied.
+    original_batch : torch.Tensor
+        Original images.
+    masks : torch.Tensor
+        Coordinates of changed pixels.
+    *args : Any
+        Additional arguments.
+
+    Returns
+    -------
+    torch.Tensor
+        Loss value.
+    """
+    errors = (original_batch - manipulated_batch) ** 2
+    # Average over pixels and batch
+    loss = torch.sum(errors * masks) / torch.sum(masks)
+    return loss  # TODO change output to dict ?
+
+
+def pn2v_loss(
+    samples: torch.Tensor,
+    labels: torch.Tensor,
+    masks: torch.Tensor,
+    noise_model: GaussianMixtureNoiseModel,
+) -> torch.Tensor:
+    """
+    Probabilistic N2V loss function described in A Krull et al., CVF (2019).
+
+    Parameters
+    ----------
+    samples : torch.Tensor # TODO this naming is confusing
+        Predicted pixel values from the network.
+    labels : torch.Tensor
+        Original pixel values.
+    masks : torch.Tensor
+        Coordinates of manipulated pixels.
+    noise_model : GaussianMixtureNoiseModel
+        Noise model for computing likelihood.
+
+    Returns
+    -------
+    torch.Tensor
+        Loss value.
+    """
+    likelihoods = noise_model.likelihood(labels, samples)
+    likelihoods_avg = torch.log(torch.mean(likelihoods, dim=1, keepdim=True))
+
+    # Average over pixels and batch
+    loss = -torch.sum(likelihoods_avg * masks) / torch.sum(masks)
+    return loss
+
+
+def mae_loss(samples: torch.Tensor, labels: torch.Tensor, *args) -> torch.Tensor:
+    """
+    N2N Loss function described in to J Lehtinen et al 2018.
+
+    Parameters
+    ----------
+    samples : torch.Tensor
+        Raw patches.
+    labels : torch.Tensor
+        Different subset of noisy patches.
+    *args : Any
+        Additional arguments.
+
+    Returns
+    -------
+    torch.Tensor
+        Loss value.
+    """
+    loss = L1Loss()
+    return loss(samples, labels)
+
+
+# def dice_loss(
+#     samples: torch.Tensor, labels: torch.Tensor, mode: str = "multiclass"
+# ) -> torch.Tensor:
+#     """Dice loss function."""
+#     return DiceLoss(mode=mode)(samples, labels.long())

careamics/losses/loss_factory.py

@@ -0,0 +1,80 @@
+"""
+Loss factory module.
+
+This module contains a factory function for creating loss functions.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Union
+
+from torch import Tensor as tensor
+
+from ..config.support import SupportedLoss
+from .fcn.losses import mae_loss, mse_loss, n2v_loss, pn2v_loss
+from .lvae.losses import (
+    denoisplit_loss,
+    denoisplit_musplit_loss,
+    hdn_loss,
+    musplit_loss,
+)
+
+
+@dataclass
+class FCNLossParameters:
+    """Dataclass for FCN loss."""
+
+    # TODO check
+    prediction: tensor
+    targets: tensor
+    mask: tensor
+    current_epoch: int
+    loss_weight: float
+
+
+def loss_factory(loss: Union[SupportedLoss, str]) -> Callable:
+    """Return loss function.
+
+    Parameters
+    ----------
+    loss : Union[SupportedLoss, str]
+        Requested loss.
+
+    Returns
+    -------
+    Callable
+        Loss function.
+
+    Raises
+    ------
+    NotImplementedError
+        If the loss is unknown.
+    """
+    if loss == SupportedLoss.N2V:
+        return n2v_loss
+
+    elif loss == SupportedLoss.PN2V:
+        return pn2v_loss
+
+    elif loss == SupportedLoss.MAE:
+        return mae_loss
+
+    elif loss == SupportedLoss.MSE:
+        return mse_loss
+
+    elif loss == SupportedLoss.HDN:
+        return hdn_loss
+
+    elif loss == SupportedLoss.MUSPLIT:
+        return musplit_loss
+
+    elif loss == SupportedLoss.DENOISPLIT:
+        return denoisplit_loss
+
+    elif loss == SupportedLoss.DENOISPLIT_MUSPLIT:
+        return denoisplit_musplit_loss
+
+    else:
+        raise NotImplementedError(f"Loss {loss} is not yet supported.")

careamics/losses/lvae/__init__.py

@@ -0,0 +1 @@
+"""LVAE losses."""

careamics/losses/lvae/loss_utils.py

@@ -0,0 +1,83 @@
+import torch
+
+
+def free_bits_kl(
+    kl: torch.Tensor, free_bits: float, batch_average: bool = False, eps: float = 1e-6
+) -> torch.Tensor:
+    """Compute free-bits version of KL divergence.
+
+    This function ensures that the KL doesn't go to zero for any latent dimension.
+    Hence, it contributes to use latent variables more efficiently, leading to
+    better representation learning.
+
+    NOTE:
+    Takes in the KL with shape (batch size, layers), returns the KL with
+    free bits (for optimization) with shape (layers,), which is the average
+    free-bits KL per layer in the current batch.
+    If batch_average is False (default), the free bits are per layer and
+    per batch element. Otherwise, the free bits are still per layer, but
+    are assigned on average to the whole batch. In both cases, the batch
+    average is returned, so it's simply a matter of doing mean(clamp(KL))
+    or clamp(mean(KL)).
+
+    Parameters
+    ----------
+    kl : torch.Tensor
+        The KL divergence tensor with shape (batch size, layers).
+    free_bits : float
+        The free bits value. Set to 0.0 to disable free bits.
+    batch_average : bool
+        Whether to average over the batch before clamping to `free_bits`.
+    eps : float
+        A small value to avoid numerical instability.
+
+    Returns
+    -------
+    torch.Tensor
+        The free-bits version of the KL divergence with shape (layers,).
+    """
+    assert kl.dim() == 2
+    if free_bits < eps:
+        return kl.mean(0)
+    if batch_average:
+        return kl.mean(0).clamp(min=free_bits)
+    return kl.clamp(min=free_bits).mean(0)
+
+
+def get_kl_weight(
+    kl_annealing: bool,
+    kl_start: int,
+    kl_annealtime: int,
+    kl_weight: float,
+    current_epoch: int,
+) -> float:
+    """Compute the weight of the KL loss in case of annealing.
+
+    Parameters
+    ----------
+    kl_annealing : bool
+        Whether to use KL annealing.
+    kl_start : int
+        The epoch at which to start
+    kl_annealtime : int
+        The number of epochs for which annealing is applied.
+    kl_weight : float
+        The weight for the KL loss. If `None`, the weight is computed
+        using annealing, else it is set to a default of 1.
+    current_epoch : int
+        The current epoch.
+    """
+    if kl_annealing:
+        # calculate relative weight
+        kl_weight = (current_epoch - kl_start) * (1.0 / kl_annealtime)
+        # clamp to [0,1]
+        kl_weight = min(max(0.0, kl_weight), 1.0)
+
+        # if the final weight is given, then apply that weight on top of it
+        if kl_weight is not None:
+            kl_weight = kl_weight * kl_weight
+    elif kl_weight is not None:
+        return kl_weight
+    else:
+        kl_weight = 1.0
+    return kl_weight