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

docs/source/losses.rst

@@ -139,6 +139,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

@@ -41,5 +41,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,200 @@
+# 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
+
+from typing import Callable, Optional
+
+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())
+ if not y.is_complex():
+ y = torch.view_as_complex(y.permute(0, 2, 3, 1).contiguous())
+
+ diff = torch.abs(x - y)
+ return nn.functional.mse_loss(diff, torch.zeros_like(diff), reduction="mean")
+
+
+def sure_loss_function(
+ operator: Callable,
+ x: torch.Tensor,
+ y_pseudo_gt: torch.Tensor,
+ y_ref: Optional[torch.Tensor] = None,
+ eps: Optional[float] = -1.0,
+ perturb_noise: Optional[torch.Tensor] = None,
+ complex_input: Optional[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. Set to -1 to set it
+ automatically estimated based on y_pseudo_gtk
+
+ 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 == -1.0:
+ eps = float(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.0 / eps * torch.matmul(perturb_noise.permute(0, 1, 3, 2), y_perturbed - y_ref)) # type: ignore
+ # 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/guide 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: Optional[torch.Tensor] = None, eps: Optional[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: Optional[torch.Tensor] = None,
+ complex_input: Optional[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.
+ """
+
+ # check inputs shapes
+ if x.dim() != 4:
+ raise ValueError(f"Input tensor x should be 4D, got {x.dim()}.")
+ if y_pseudo_gt.dim() != 4:
+ raise ValueError(f"Input tensor y_pseudo_gt should be 4D, but got {y_pseudo_gt.dim()}.")
+ if y_ref is not None and y_ref.dim() != 4:
+ raise ValueError(f"Input tensor y_ref should be 4D, but got {y_ref.dim()}.")
+ if x.shape != y_pseudo_gt.shape:
+ raise ValueError(
+ f"Input tensor x and y_pseudo_gt should have the same shape, but got x shape {x.shape}, "
+ f"y_pseudo_gt shape {y_pseudo_gt.shape}."
+ )
+ if y_ref is not None and y_pseudo_gt.shape != y_ref.shape:
+ raise ValueError(
+ f"Input tensor y_pseudo_gt and y_ref should have the same shape, but got y_pseudo_gt shape {y_pseudo_gt.shape}, "
+ f"y_ref shape {y_ref.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,112 @@
+# 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
+
+from typing import Callable
+
+import torch
+from torch import nn
+
+
+def _zdot(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+ """
+ Complex dot product between tensors x1 and x2: sum(x1.*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(x: torch.Tensor) -> torch.Tensor:
+ """
+ Complex dot product between tensor x and itself
+ """
+ res = _zdot(x, x)
+ if torch.is_complex(res):
+ return res.real
+ else:
+ return res
+
+
+class ConjugateGradient(nn.Module):
+ """
+ Congugate Gradient (CG) solver for linear systems Ax = y.
+
+ For 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):
+ """
+ Args:
+ linear_op: Linear operator
+ num_iter: Number of iterations to run CG
+ """
+ super().__init__()
+
+ self.linear_op = linear_op
+ self.num_iter = num_iter
+
+ def update(
+ self, x: torch.Tensor, p: torch.Tensor, r: torch.Tensor, rsold: torch.Tensor
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ 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.
+ """
+
+ dy = self.linear_op(p)
+ p_dot_dy = _zdot(p, dy)
+ alpha = rsold / p_dot_dy
+ x = x + alpha * p
+ r = r - alpha * dy
+ rsnew = _zdot_single(r)
+ beta = rsnew / rsold
+ rsold = rsnew
+ p = beta * p + r
+ return x, p, r, rsold
+
+ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+ """
+ run conjugate gradient for num_iter iterations to solve Ax = y
+
+ Args:
+ x: tensor (real or complex); Initial guess for linear system Ax = y.
+ The size of x should be applicable to the linear operator. For
+ example, if the linear operator is FFT, then x is HCHW; if the
+ linear operator is a matrix multiplication, then x is a vector
+
+ y: tensor (real or complex); Measurement. Same size as x
+
+ Returns:
+ x: Solution to Ax = y
+ """
+ # Compute residual
+ r = y - self.linear_op(x)
+ rsold = _zdot_single(r)
+ p = r
+
+ # Update
+ for _i in range(self.num_iter):
+ x, p, r, rsold = self.update(x, p, r, rsold)
+ if rsold < 1e-10:
+ break
+ return x