deepinv 0.1.0.dev0__py3-none-any.whl

deepinv/loss/metric.py

@@ -0,0 +1,125 @@
+import math
+import torch
+import torch.nn as nn
+from torch import autograd as autograd
+
+
+class LpNorm(torch.nn.Module):
+    r"""
+    :math:`\ell_p` metric for :math:`p>0`.
+
+
+    If ``onesided=False`` then the metric is defined as
+    :math:`d(x,y)=\|x-y\|_p^p`.
+
+    otherwise it is the one-sided error https://ieeexplore.ieee.org/abstract/document/6418031/, defined as
+    :math:`d(x,y)= \|\max(x\circ y) \|_p^p`. where :math:`\circ` denotes element-wise multiplication.
+
+    """
+
+    def __init__(self, p=2, onesided=False):
+        super().__init__()
+        self.p = p
+        self.onesided = onesided
+
+    def forward(self, x, y):
+        if self.onesided:
+            return torch.nn.functional.relu(-x * y).flatten().pow(self.p).mean()
+        else:
+            return (x - y).flatten().abs().pow(self.p).mean()
+
+
+def mse():
+    return nn.MSELoss()
+
+
+def l1():
+    return nn.L1Loss()
+
+
+class CharbonnierLoss(nn.Module):
+    r"""
+    Charbonnier Loss
+
+    """
+
+    def __init__(self, eps=1e-9):
+        super(CharbonnierLoss, self).__init__()
+        self.eps = eps
+
+    def forward(self, x, y):
+        diff = x - y
+        loss = torch.mean(torch.sqrt((diff * diff) + self.eps))
+        return loss
+
+
+def r1_penalty(real_pred, real_img):
+    """R1 regularization for discriminator. The core idea is to
+    penalize the gradient on real data alone: when the
+    generator distribution produces the true data distribution
+    and the discriminator is equal to 0 on the data manifold, the
+    gradient penalty ensures that the discriminator cannot create
+    a non-zero gradient orthogonal to the data manifold without
+    suffering a loss in the GAN game.
+    Ref:
+    Eq. 9 in Which training methods for GANs do actually converge.
+    """
+    grad_real = autograd.grad(
+        outputs=real_pred.sum(), inputs=real_img, create_graph=True
+    )[0]
+    grad_penalty = grad_real.pow(2).view(grad_real.shape[0], -1).sum(1).mean()
+    return grad_penalty
+
+
+def g_path_regularize(fake_img, latents, mean_path_length, decay=0.01):
+    noise = torch.randn_like(fake_img) / math.sqrt(
+        fake_img.shape[2] * fake_img.shape[3]
+    )
+    grad = autograd.grad(
+        outputs=(fake_img * noise).sum(), inputs=latents, create_graph=True
+    )[0]
+    path_lengths = torch.sqrt(grad.pow(2).sum(2).mean(1))
+
+    path_mean = mean_path_length + decay * (path_lengths.mean() - mean_path_length)
+
+    path_penalty = (path_lengths - path_mean).pow(2).mean()
+
+    return path_penalty, path_lengths.detach().mean(), path_mean.detach()
+
+
+def gradient_penalty_loss(discriminator, real_data, fake_data, weight=None):
+    """Calculate gradient penalty for wgan-gp.
+    Args:
+        discriminator (nn.Module): Network for the discriminator.
+        real_data (Tensor): Real input data.
+        fake_data (Tensor): Fake input data.
+        weight (Tensor): Weight tensor. Default: None.
+    Returns:
+        Tensor: A tensor for gradient penalty.
+    """
+
+    batch_size = real_data.size(0)
+    alpha = real_data.new_tensor(torch.rand(batch_size, 1, 1, 1))
+
+    # interpolate between real_data and fake_data
+    interpolates = alpha * real_data + (1.0 - alpha) * fake_data
+    interpolates = autograd.Variable(interpolates, requires_grad=True)
+
+    disc_interpolates = discriminator(interpolates)
+    gradients = autograd.grad(
+        outputs=disc_interpolates,
+        inputs=interpolates,
+        grad_outputs=torch.ones_like(disc_interpolates),
+        create_graph=True,
+        retain_graph=True,
+        only_inputs=True,
+    )[0]
+
+    if weight is not None:
+        gradients = gradients * weight
+
+    gradients_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
+    if weight is not None:
+        gradients_penalty /= torch.mean(weight)
+
+    return gradients_penalty

deepinv/loss/moi.py

@@ -0,0 +1,64 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class MOILoss(nn.Module):
+    r"""
+    Multi-operator imaging loss
+
+    This loss can be used to learn when signals are observed via multiple (possibly incomplete)
+    forward operators :math:`\{A_g\}_{g=1}^{G}`,
+    i.e., :math:`y_i = A_{g_i}x_i` where :math:`g_i\in \{1,\dots,G\}` (see https://arxiv.org/abs/2201.12151).
+
+
+    The measurement consistency loss is defined as
+
+    .. math::
+
+        \| \hat{x} - \inverse{\hat{x},A_g} \|^2
+
+    where :math:`\hat{x}=\inverse{y,A_s}` is a reconstructed signal (observed via operator :math:`A_s`) and
+    :math:`A_g` is a forward operator sampled at random from a set :math:`\{A_g\}_{g=1}^{G}`.
+
+    By default, the error is computed using the MSE metric, however any other metric (e.g., :math:`\ell_1`)
+    can be used as well.
+
+    :param torch.nn.Module metric: metric used for computing data consistency,
+        which is set as the mean squared error by default.
+    :param float weight: total weight of the loss
+    :param bool apply_noise: if ``True``, the augmented measurement is computed with the full sensing model
+        :math:`\sensor{\noise{\forw{\hat{x}}}}` (i.e., noise and sensor model),
+        otherwise is generated as :math:`\forw{\hat{x}}`.
+    """
+
+    def __init__(
+        self, physics_list, metric=torch.nn.MSELoss(), apply_noise=True, weight=1.0
+    ):
+        super(MOILoss, self).__init__()
+        self.name = "moi"
+        self.physics_list = physics_list
+        self.metric = metric
+        self.weight = weight
+        self.noise = apply_noise
+
+    def forward(self, x_net, model, **kwargs):
+        r"""
+        Computes the MOI loss.
+
+        :param torch.Tensor x_net: Reconstructed image :math:`\inverse{y}`.
+        :param list of deepinv.physics.Physics physics: List containing the :math:`G` different forward operators
+            associated with the measurements.
+        :param torch.nn.Module model: Reconstruction function.
+        :return: (torch.Tensor) loss.
+        """
+        j = np.random.randint(len(self.physics_list))
+
+        if self.noise:
+            y = self.physics_list[j](x_net)
+        else:
+            y = self.physics_list[j].A(x_net)
+
+        x2 = model(y, self.physics_list[j])
+
+        return self.weight * self.metric(x2, x_net)

deepinv/loss/regularisers.py

@@ -0,0 +1,155 @@
+import torch
+import torch.nn as nn
+
+
+class JacobianSpectralNorm(nn.Module):
+    r"""
+    Computes the spectral norm of the Jacobian.
+
+    Given a function :math:`f:\mathbb{R}^n\to\mathbb{R}^n`, this module computes the spectral
+    norm of the Jacobian of :math:`f` in :math:`x`, i.e.
+
+    .. math::
+        \|\frac{df}{du}(x)\|_2.
+
+    This spectral norm is computed with a power method leveraging jacobian vector products, as proposed in `<https://arxiv.org/abs/2012.13247v2>`_.
+
+    :param int max_iter: maximum numer of iteration of the power method.
+    :param float tol: tolerance for the convergence of the power method.
+    :param bool eval_mode: set to `False` if one does not want to backpropagate through the spectral norm (default), set to `True` otherwise.
+    :param bool verbose: whether to print computation details or not.
+
+
+    Example of usage:
+
+    ::
+
+        import torch
+        from deepinv.loss.regularisers import JacobianSpectralNorm
+
+        reg_l2 = JacobianSpectralNorm(max_iter=10, tol=1e-3, eval_mode=False, verbose=True)
+
+        A = torch.diag(torch.Tensor(range(1, 51)))  # creates a diagonal matrix with largest eigenvalue = 50
+        x = torch.randn_like(A).requires_grad_()
+        out = A @ x
+
+        regval = reg_l2(out, x)
+        print(regval) # >> returns approx 50
+    """
+
+    def __init__(self, max_iter=10, tol=1e-3, eval_mode=False, verbose=False):
+        super(JacobianSpectralNorm, self).__init__()
+        self.name = "jsn"
+        self.max_iter = max_iter
+        self.tol = tol
+        self.eval = eval_mode
+        self.verbose = verbose
+
+    def forward(self, y, x, **kwargs):
+        """
+        Computes the spectral norm of the Jacobian of :math:`f` in :math:`x`.
+
+        .. warning::
+            The input :math:`x` must have requires_grad=True before evaluating :math:`f`.
+
+        :param torch.Tensor y: output of the function :math:`f` at :math:`x`.
+        :param torch.Tensor x: input of the function :math:`f`.
+        """
+        u = torch.randn_like(x)
+        u = u / torch.norm(u.flatten(), p=2)
+
+        zold = torch.zeros_like(u)
+
+        for it in range(self.max_iter):
+            # Double backward trick. From https://gist.github.com/apaszke/c7257ac04cb8debb82221764f6d117ad
+            w = torch.ones_like(y, requires_grad=True)
+            v = torch.autograd.grad(
+                torch.autograd.grad(y, x, w, create_graph=True),
+                w,
+                u,
+                create_graph=not self.eval,
+            )[
+                0
+            ]  # v = A(u)
+
+            (v,) = torch.autograd.grad(y, x, v, retain_graph=True, create_graph=True)
+
+            z = torch.dot(u.flatten(), v.flatten()) / torch.norm(u, p=2) ** 2
+
+            if it > 0:
+                rel_var = torch.norm(z - zold)
+                if rel_var < self.tol and self.verbose:
+                    print(
+                        "Power iteration converged at iteration: ",
+                        it,
+                        ", val: ",
+                        z.sqrt().item(),
+                        ", relvar :",
+                        rel_var,
+                    )
+                    break
+            zold = z.detach().clone()
+
+            u = v / torch.norm(v.flatten(), p=2)
+
+            if self.eval:
+                w.detach_()
+                v.detach_()
+                u.detach_()
+
+        return z.view(-1).sqrt()
+
+
+class FNEJacobianSpectralNorm(nn.Module):
+    r"""
+    Computes the Firm-Nonexpansiveness Jacobian spectral norm.
+
+    Given a function :math:`f:\mathbb{R}^n\to\mathbb{R}^n`, this module computes the spectral
+    norm of the Jacobian of :math:`2f-\operatorname{Id}` (where :math:`\operatorname{Id}` denotes the
+    identity) in :math:`x`, i.e.
+
+    .. math::
+        \|\frac{d(2f-\operatorname{Id})}{du}(x)\|_2,
+
+    as proposed in `<https://arxiv.org/abs/2012.13247v2>`_.
+    This spectral norm is computed with the :meth:`deepinv.loss.JacobianSpectralNorm` module.
+
+    :param int max_iter: maximum numer of iteration of the power method.
+    :param float tol: tolerance for the convergence of the power method.
+    :param bool eval_mode: set to `False` if one does not want to backpropagate through the spectral norm (default), set to `True` otherwise.
+    :param bool verbose: whether to print computation details or not.
+
+    """
+
+    def __init__(self, max_iter=10, tol=1e-3, verbose=False, eval_mode=False):
+        super(FNEJacobianSpectralNorm, self).__init__()
+        self.spectral_norm_module = JacobianSpectralNorm(
+            max_iter=max_iter, tol=tol, verbose=verbose, eval_mode=eval_mode
+        )
+
+    def forward(
+        self, y_in, x_in, model, *args_model, interpolation=False, **kwargs_model
+    ):
+        r"""
+        Computes the Firm-Nonexpansiveness (FNE) Jacobian spectral norm of a model.
+
+        :param torch.Tensor y_in: input of the model (by default).
+        :param torch.Tensor x_in: an additional point of the model (by default).
+        :param torch.nn.Module model: neural network, or function, of which we want to compute the FNE Jacobian spectral norm.
+        :param `*args_model`: additional arguments of the model.
+        :param bool interpolation: whether to input to model an interpolation between y_in and x_in instead of y_in (default is `False`).
+        :param `**kargs_model`: additional keyword arguments of the model.
+        """
+
+        if interpolation:
+            eta = torch.rand(y_in.size(0), 1, 1, 1, requires_grad=True).to(y_in.device)
+            x = eta * y_in.detach() + (1 - eta) * x_in.detach()
+        else:
+            x = y_in
+
+        x.requires_grad_()
+        x_out = model(x, *args_model, **kwargs_model)
+
+        y = 2.0 * x_out - x
+
+        return self.spectral_norm_module(y, x)

deepinv/loss/score.py

@@ -0,0 +1,41 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class ScoreLoss(nn.Module):
+    r"""
+    Approximates score of a distribution.
+
+    It approximates the score of the measurement distribution :math:`S(y)\approx \nabla \log p(y)`
+    https://proceedings.neurips.cc/paper_files/paper/2021/file/077b83af57538aa183971a2fe0971ec1-Paper.pdf.
+
+    The score loss is defined as
+
+    .. math::
+
+        \| \epsilon - \sigma S(y+ \sigma \epsilon) \|^2
+
+    where :math:`\epsilon` is sampled from :math:`N(0,I)` and
+    :math:`\sigma` is sampled from :math:`N(0,I\delta^2)`.
+
+    :param float delta: hyperparameter :math:`\delta` controlling the level of noise.
+    """
+
+    def __init__(self, delta):
+        super(ScoreLoss, self).__init__()
+        self.name = "score"
+        self.metric = torch.nn.MSELoss()
+        self.delta = delta
+
+    def forward(self, y, model, **kwargs):
+        r"""
+        Computes the Score loss.
+
+        :param torch.Tensor y: measurements.
+        :param torch.nn.Module model: Reconstruction function.
+        :return: (torch.Tensor) loss.
+        """
+        std = np.randn() * self.delta
+        noise = torch.randn_like(y)
+        return self.metric(noise, std * model(std * noise + y))

deepinv/loss/sup.py

@@ -0,0 +1,37 @@
+import torch
+import torch.nn as nn
+
+
+class SupLoss(nn.Module):
+    r"""
+    Standard supervised loss
+
+    The supervised loss is defined as
+
+    .. math::
+
+        \|x-\inverse{y}\|^2
+
+    where :math:`\inverse{y}` is the reconstructed signal and :math:`x` is the ground truth target.
+
+    By default, the error is computed using the MSE metric, however any other metric (e.g., :math:`\ell_1`)
+    can be used as well.
+
+    :param torch.nn.Module metric: metric used for computing data consistency,
+        which is set as the mean squared error by default.
+    """
+
+    def __init__(self, metric=torch.nn.MSELoss()):
+        super(SupLoss, self).__init__()
+        self.name = "supervised"
+        self.metric = metric
+
+    def forward(self, x, x_net, **kwargs):
+        r"""
+        Computes the loss.
+
+        :param torch.Tensor x: Target (ground-truth) image.
+        :param torch.Tensor x_net: Reconstructed image :math:\inverse{y}.
+        :return: (torch.Tensor) loss.
+        """
+        return self.metric(x_net, x)