Stein's Unbiased Risk Estimator (SURE) loss and Conjugate Gradient

docs/source/losses.rst

@@ -134,6 +134,11 @@ Reconstruction Losses
 .. autoclass:: JukeboxLoss
     :members:
 
+`SURELoss`
+~~~~~~~~~~~~~~
+.. autoclass:: SURELoss
+    :members:
+
 
 Loss Wrappers
 -------------

docs/source/networks.rst

@@ -408,6 +408,11 @@ Layers
 .. autoclass:: LLTM
     :members:
 
+`ConjugateGradient`
+~~~~~~~~~~~~~~~~~~~
+.. autoclass:: ConjugateGradient
+    :members:
+
 `Utilities`
 ~~~~~~~~~~~
 .. automodule:: monai.networks.layers.convutils

monai/losses/__init__.py

@@ -40,5 +40,6 @@
 from .spatial_mask import MaskedLoss
 from .spectral_loss import JukeboxLoss
 from .ssim_loss import SSIMLoss
+from .sure_loss import SURELoss
 from .tversky import TverskyLoss
 from .unified_focal_loss import AsymmetricUnifiedFocalLoss

monai/losses/sure_loss.py

@@ -0,0 +1,198 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch.nn.modules.loss import _Loss
+
+
+def complex_diff_abs_loss(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    """
+    First compute the difference in the complex domain,
+    then get the absolute value and take the mse
+
+    Args:
+        x, y - B, 2, H, W real valued tensors representing complex numbers
+                or  B,1,H,W complex valued tensors
+    Returns:
+        l2_loss - scalar
+    """
+    if not x.is_complex():
+        x_ = torch.view_as_complex(x.permute(0, 2, 3, 1).contiguous())
+    else:
+        x_ = x
+    if not y.is_complex():
+        y_ = torch.view_as_complex(y.permute(0, 2, 3, 1).contiguous())
+    else:
+        y_ = y
+    diff = x_ - y_
+    l2_loss = nn.functional.mse_loss(torch.abs(diff), torch.zeros_like(torch.abs(diff)), reduction="mean")
+    return l2_loss
+
+
+def sure_loss_function(
+    operator: callable,
+    x: torch.Tensor,
+    y_pseudo_gt: torch.Tensor,
+    y_ref: torch.Tensor = None,
+    eps: float = None,
+    perturb_noise: torch.Tensor = None,
+    complex_input: bool = False,
+) -> torch.Tensor:
+    """
+
+    Args:
+        operator (function): The operator function that takes in an input
+        tensor x and returns an output tensor y. We will use this to compute
+        the divergence. More specifically, we will perturb the input x by a
+        small amount and compute the divergence between the perturbed output
+        and the reference output
+
+        x (torch.Tensor): The input tensor of shape (B, C, H, W) to the
+        operator.  For complex input, the shape is (B, 2, H, W) aka C=2 real.
+        For real input, the shape is (B, 1, H, W) real.
+
+        y_pseudo_gt (torch.Tensor): The pseudo ground truth tensor of shape (B,
+        C, H, W) used to compute the L2 loss.  For complex input, the shape is
+        (B, 2, H, W) aka C=2 real.  For real input, the shape is (B, 1, H, W)
+        real.
+
+        y_ref (torch.Tensor, optional): The reference output tensor of shape
+        (B, C, H, W) used to compute the divergence. Defaults to None.  For
+        complex input, the shape is (B, 2, H, W) aka C=2 real.  For real input,
+        the shape is (B, 1, H, W) real.
+
+        eps (float, optional): The perturbation scalar. Defaults to None.
+
+        perturb_noise (torch.Tensor, optional): The noise vector of shape (B,
+        C, H, W). Defaults to None.  For complex input, the shape is (B, 2, H,
+        W) aka C=2 real.  For real input, the shape is (B, 1, H, W) real.
+
+        complex_input(bool, optional): Whether the input is complex or not.
+        Defaults to False.
+
+    Returns:
+        sure_loss (torch.Tensor): The SURE loss scalar.
+    """
+    # perturb input
+    if perturb_noise is None:
+        perturb_noise = torch.randn_like(x)
+    if eps is None:
+        eps = torch.abs(y_pseudo_gt.max()) / 1000
+    # get y_ref if not provided
+    if y_ref is None:
+        y_ref = operator(x)
+
+    # get perturbed output
+    x_perturbed = x + eps * perturb_noise
+    y_perturbed = operator(x_perturbed)
+    # divergence
+    divergence = torch.sum(1 / eps * torch.matmul(perturb_noise.permute(0, 1, 3, 2), y_perturbed - y_ref))
+    # l2 loss between y_ref, y_pseudo_gt
+    if complex_input:
+        l2_loss = complex_diff_abs_loss(y_ref, y_pseudo_gt)
+    else:
+        # real input
+        l2_loss = nn.functional.mse_loss(y_ref, y_pseudo_gt, reduction="mean")
+
+    # sure loss
+    sure_loss = l2_loss * divergence / (x.shape[0] * x.shape[2] * x.shape[3])
+    return sure_loss
+
+
+class SURELoss(_Loss):
+    """
+    Calculate the Stein's Unbiased Risk Estimator (SURE) loss for a given operator.
+
+    This is a differentiable loss function that can be used to train/giude an
+    operator (e.g. neural network), where the pseudo ground truth is available
+    but the reference ground truth is not. For example, in the MRI
+    reconstruction, the pseudo ground truth is the zero-filled reconstruction
+    and the reference ground truth is the fully sampled reconstruction.  Often,
+    the reference ground truth is not available due to the lack of fully sampled
+    data.
+
+    The original SURE loss is proposed in [1]. The SURE loss used for guiding
+    the diffusion model based MRI reconstruction is proposed in [2].
+
+    Reference
+
+    [1] Stein, C.M.: Estimation of the mean of a multivariate normal distribution. Annals of Statistics
+
+    [2] B. Ozturkler et al. SMRD: SURE-based Robust MRI Reconstruction with Diffusion Models.
+    (https://arxiv.org/pdf/2310.01799.pdf)
+    """
+
+    def __init__(self, perturb_noise: torch.Tensor = None, eps: float = None) -> None:
+        """
+        Args:
+            perturb_noise (torch.Tensor, optional): The noise vector of shape (B, C, H, W). Defaults to None.  For complex input, the shape is (B, 2, H, W) aka C=2 real.  For real input, the shape is (B, 1, H, W) real.
+            eps (float, optional): The perturbation scalar. Defaults to None.
+        """
+
+        super().__init__()
+        self.perturb_noise = perturb_noise
+        self.eps = eps
+
+    def forward(
+        self,
+        operator: callable,
+        x: torch.Tensor,
+        y_pseudo_gt: torch.Tensor,
+        y_ref: torch.Tensor = None,
+        complex_input: bool = False,
+    ) -> torch.Tensor:
+        """
+        Args:
+            operator (function): The operator function that takes in an input tensor x and returns an output tensor y. We will use this to compute the divergence. More specifically, we will perturb the input x by a small amount and compute the divergence between the perturbed output and the reference output
+            x (torch.Tensor): The input tensor of shape (B, C, H, W) to the operator. C=1 or 2: For complex input, the shape is (B, 2, H, W) aka C=2 real.  For real input, the shape is (B, 1, H, W) real.
+            y_pseudo_gt (torch.Tensor): The pseudo ground truth tensor of shape (B, C, H, W) used to compute the L2 loss. C=1 or 2: For complex input, the shape is (B, 2, H, W) aka C=2 real.  For real input, the shape is (B, 1, H, W) real.
+            y_ref (torch.Tensor, optional): The reference output tensor of the same shape as y_pseudo_gt
+
+        Returns:
+            sure_loss (torch.Tensor): The SURE loss scalar.
+        """
+        # TODO: support for C>2 real valued input
+
+        # check inputs
+        # dim check:
+        assert x.dim() == 4, "Input tensor x should be 4D."
+        assert y_pseudo_gt.dim() == 4, "Input tensor y_pseudo_gt should be 4D."
+        if y_ref is not None:
+            assert y_ref.dim() == 4, "Input tensor y_ref should be 4D."
+
+        # complex/real check:
+        if complex_input:
+            assert (
+                x.shape[1] == 2 and y_pseudo_gt.shape[1] == 2 and not x.is_complex() and not y_pseudo_gt.is_complex()
+            ), "For complex input, the shape is (B, 2, H, W) aka C=2 real or (B, 1, H, W) aka C=1 complex"
+            if y_ref is not None:
+                assert (
+                    y_ref.shape[1] == 2 and not y_ref.is_complex()
+                ), "For complex input, the shape is (B, 2, H, W) aka C=2 real or (B, 1, H, W) aka C=1 complex"
+        else:  # real input
+            assert (
+                x.shape[1] == 1 and y_pseudo_gt.shape[1] == 1 and not x.is_complex() and not y_pseudo_gt.is_complex()
+            ), "For real input, the shape is (B, 1, H, W) real."
+            if y_ref is not None:
+                assert y_ref.shape[1] == 1 and not y_ref.is_complex(), "For real input, the shape is (B, 1, H, W) real."
+
+        # shape check
+        assert x.shape == y_pseudo_gt.shape, "Input tensor x and y_pseudo_gt should have the same shape."
+        if y_ref is not None:
+            assert y_pseudo_gt.shape == y_ref.shape, "Input tensor y_pseudo_gt and y_ref should have the same shape."
+
+        # compute loss
+        loss = sure_loss_function(operator, x, y_pseudo_gt, y_ref, self.eps, self.perturb_noise, complex_input)
+
+        return loss

monai/networks/layers/__init__.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+from .conjugate_gradient import ConjugateGradient
 from .convutils import calculate_out_shape, gaussian_1d, polyval, same_padding, stride_minus_kernel_padding
 from .drop_path import DropPath
 from .factories import Act, Conv, Dropout, LayerFactory, Norm, Pad, Pool, split_args

monai/networks/layers/conjugate_gradient.py

@@ -0,0 +1,114 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class ConjugateGradient(nn.Module):
+    """
+    Congugate Gradient (CG) solver for linear systems Ax = y.
+
+    For A (linear_op) that is positive definite and self-adjoint, CG is
+    guaranteed to converge CG is often used to solve linear systems of the form
+    Ax = y, where A is too large to store explicitly, but can be computed via a
+    linear operator.
+
+    As a result, here we won't set A explicitly as a matrix, but rather as a
+    linear operator. For example, A could be a FFT/IFFT operation
+    """
+
+    def __init__(self, linear_op: callable, num_iter: int, dbprint: bool = False):
+        """
+        Args:
+            linear_op: Linear operator
+            num_iter: Number of iterations to run CG
+            dbprint [False]: Print residual at each iteration
+        """
+        super().__init__()
+
+        self.A = linear_op
+        self.num_iter = num_iter
+        self.dbprint = dbprint
+
+    def _zdot(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+        """
+        Complex dot product between tensors x1 and x2.
+        """
+        if torch.is_complex(x1):
+            assert torch.is_complex(x2), "x1 and x2 must both be complex"
+            return torch.sum(x1.conj() * x2)
+        else:
+            return torch.sum(x1 * x2)
+
+    def _zdot_single(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Complex dot product between tensor x and itself
+        """
+        res = self._zdot(x, x)
+        if torch.is_complex(res):
+            return res.real
+        else:
+            return res
+
+    def _update(self, iter: int) -> callable:
+        """
+        perform one iteration of the CG method. It takes the current solution x,
+        the current search direction p, the current residual r, and the old
+        residual norm rsold as inputs. Then it computes the new solution, search
+        direction, residual, and residual norm, and returns them.
+        """
+
+        def update_fn(
+            x: torch.Tensor, p: torch.Tensor, r: torch.Tensor, rsold: torch.Tensor
+        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+            dy = self.A(p)
+            p_dot_dy = self._zdot(p, dy)
+            alpha = rsold / p_dot_dy
+            x = x + alpha * p
+            r = r - alpha * dy
+            rsnew = self._zdot_single(r)
+            beta = rsnew / rsold
+            rsold = rsnew
+            p = beta * p + r
+
+            # print residual
+            if self.dbprint:
+                print(f"CG Iteration {iter}: {rsnew}")
+
+            return x, p, r, rsold
+
+        return update_fn
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        """
+        run conjugate gradient for num_iter iterations to solve Ax = y
+
+        Args:
+            x: B H W tensor (real or complex); Initial guess for linear system Ax = y
+            y: B H W tensor (real or complex); Measurement
+
+        Returns:
+            x: Solution to Ax = y
+        """
+        # Compute residual
+        r = y - self.A(x)
+        rsold = self._zdot_single(r)
+        p = r
+
+        # Update
+        for i in range(self.num_iter):
+            x, p, r, rsold = self._update(i)(x, p, r, rsold)
+            if rsold < 1e-10:
+                break
+        return x