careamics 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

careamics/losses/fcn/losses.py

@@ -0,0 +1,98 @@
+"""
+Loss submodule.
+
+This submodule contains the various losses used in CAREamics.
+"""
+
+import torch
+from torch.nn import L1Loss, MSELoss
+
+
+def mse_loss(source: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    """
+    Mean squared error loss.
+
+    Parameters
+    ----------
+    source : torch.Tensor
+        Source patches.
+    target : torch.Tensor
+        Target patches.
+
+    Returns
+    -------
+    torch.Tensor
+        Loss value.
+    """
+    loss = MSELoss()
+    return loss(source, target)
+
+
+def n2v_loss(
+    manipulated_patches: torch.Tensor,
+    original_patches: torch.Tensor,
+    masks: torch.Tensor,
+) -> torch.Tensor:
+    """
+    N2V Loss function described in A Krull et al 2018.
+
+    Parameters
+    ----------
+    manipulated_patches : torch.Tensor
+        Patches with manipulated pixels.
+    original_patches : torch.Tensor
+        Noisy patches.
+    masks : torch.Tensor
+        Array containing masked pixel locations.
+
+    Returns
+    -------
+    torch.Tensor
+        Loss value.
+    """
+    errors = (original_patches - manipulated_patches) ** 2
+    # Average over pixels and batch
+    loss = torch.sum(errors * masks) / torch.sum(masks)
+    return loss  # TODO change output to dict ?
+
+
+def mae_loss(samples: torch.Tensor, labels: torch.Tensor) -> 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.
+
+    Returns
+    -------
+    torch.Tensor
+        Loss value.
+    """
+    loss = L1Loss()
+    return loss(samples, labels)
+
+
+# def pn2v_loss(
+#     samples: torch.Tensor,
+#     labels: torch.Tensor,
+#     masks: torch.Tensor,
+#     noise_model: HistogramNoiseModel,
+# ) -> torch.Tensor:
+#     """Probabilistic N2V loss function described in A Krull et al., CVF (2019)."""
+#     likelihoods = noise_model.likelihood(labels, samples)
+#     likelihoods_avg = torch.log(torch.mean(likelihoods, dim=0, keepdim=True)[0, ...])
+
+#     # Average over pixels and batch
+#     loss = -torch.sum(likelihoods_avg * masks) / torch.sum(masks)
+#     return loss
+
+
+# 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,155 @@
+"""
+Loss factory module.
+
+This module contains a factory function for creating loss functions.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, Literal, Optional, Union
+
+from torch import Tensor as tensor
+
+from ..config.support import SupportedLoss
+from .fcn.losses import mae_loss, mse_loss, n2v_loss
+from .lvae.losses import denoisplit_loss, denoisplit_musplit_loss, musplit_loss
+
+if TYPE_CHECKING:
+    from careamics.models.lvae.likelihoods import (
+        GaussianLikelihood,
+        NoiseModelLikelihood,
+    )
+    from careamics.models.lvae.noise_models import (
+        GaussianMixtureNoiseModel,
+        MultiChannelNoiseModel,
+    )
+
+    NoiseModel = Union[GaussianMixtureNoiseModel, MultiChannelNoiseModel]
+
+
+@dataclass
+class FCNLossParameters:
+    """Dataclass for FCN loss."""
+
+    # TODO check
+    prediction: tensor
+    targets: tensor
+    mask: tensor
+    current_epoch: int
+    loss_weight: float
+
+
+@dataclass  # TODO why not pydantic?
+class LVAELossParameters:
+    """Dataclass for LVAE loss."""
+
+    # TODO: refactor in more modular blocks (otherwise it gets messy very easily)
+    # e.g., - weights, - kl_params, ...
+
+    noise_model_likelihood: Optional[NoiseModelLikelihood] = None
+    """Noise model likelihood instance."""
+    gaussian_likelihood: Optional[GaussianLikelihood] = None
+    """Gaussian likelihood instance."""
+    current_epoch: int = 0
+    """Current epoch in the training loop."""
+    reconstruction_weight: float = 1.0
+    """Weight for the reconstruction loss in the total net loss
+    (i.e., `net_loss = reconstruction_weight * rec_loss + kl_weight * kl_loss`)."""
+    musplit_weight: float = 0.0
+    """Weight for the muSplit loss (used in the muSplit-deonoiSplit loss)."""
+    denoisplit_weight: float = 1.0
+    """Weight for the denoiSplit loss (used in the muSplit-deonoiSplit loss)."""
+    kl_type: Literal["kl", "kl_restricted", "kl_spatial", "kl_channelwise"] = "kl"
+    """Type of KL divergence used as KL loss."""
+    kl_weight: float = 1.0
+    """Weight for the KL loss in the total net loss.
+    (i.e., `net_loss = reconstruction_weight * rec_loss + kl_weight * kl_loss`)."""
+    kl_annealing: bool = False
+    """Whether to apply KL loss annealing."""
+    kl_start: int = -1
+    """Epoch at which KL loss annealing starts."""
+    kl_annealtime: int = 10
+    """Number of epochs for which KL loss annealing is applied."""
+    non_stochastic: bool = False
+    """Whether to sample latents and compute KL."""
+
+
+# TODO: really needed?
+# like it is now, it is difficult to use, we need a way to specify the
+# loss parameters in a more user-friendly way.
+def loss_parameters_factory(
+    type: SupportedLoss,
+) -> Union[FCNLossParameters, LVAELossParameters]:
+    """Return loss parameters.
+
+    Parameters
+    ----------
+    type : SupportedLoss
+        Requested loss.
+
+    Returns
+    -------
+    Union[FCNLossParameters, LVAELossParameters]
+        Loss parameters.
+
+    Raises
+    ------
+    NotImplementedError
+        If the loss is unknown.
+    """
+    if type in [SupportedLoss.N2V, SupportedLoss.MSE, SupportedLoss.MAE]:
+        return FCNLossParameters
+
+    elif type in [
+        SupportedLoss.MUSPLIT,
+        SupportedLoss.DENOISPLIT,
+        SupportedLoss.DENOISPLIT_MUSPLIT,
+    ]:
+        return LVAELossParameters  # it returns the class, not an instance
+
+    else:
+        raise NotImplementedError(f"Loss {type} is not yet supported.")
+
+
+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_type == SupportedLoss.PN2V:
+    #     return pn2v_loss
+
+    elif loss == SupportedLoss.MAE:
+        return mae_loss
+
+    elif loss == SupportedLoss.MSE:
+        return mse_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