Repository: Project-MONAI/MONAI
File: monai/losses/sure_loss.py (Python)

Code before:
    # 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/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: 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.


Review comment (suggestion):
```suggestion
    This is a differentiable loss function that can be used to train/guide an
```

Code after:
    # 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.