ncc_loss can not use with mask

[robert-graf]

Mask in ncc_loss from deepali.losses.functional does not work.

the image-loss is already collapses to shape=[c] before masked_loss is called, through the sum(dim=1) calls.

[qiuhuaqi]

as we discussed earlier, the keepdim flag should avoid this in the most recent version. Can you send a minimal example to reproduce the issue if this still doesn't work for you?

[robert-graf]

I figured it out; The reshape removes spatial dimensions:
source = source.reshape(source.shape[0], -1).float()
will reshape to (batch,-1)
The mean will collapse the second dimension, even if you set keepdim=True
This is intended behavoure (print(torch.zeros((10, 10, 10, 10)).mean(1, keepdim=True).shape) will be (10,1,10,10)
The shape is (batch) from that point onward.

Minimal example, shape when mask is used is 10.

ncc_loss(torch.zeros((10, 10, 10, 10)), torch.zeros((10, 10, 10, 10)), torch.zeros((10,1, 10, 10)))

def ncc_loss(
    source: Tensor,
    target: Tensor,
    mask: Optional[Tensor] = None,
    epsilon: float = 1e-15,
    reduction: str = "mean",
) -> Tensor:
    r"""Normalized cross correlation.

    Args:
        source: Source image sampled on ``target`` grid.
        target: Target image with same shape as ``source``.
        mask: Multiplicative mask tensor with same shape as ``source``.
        epsilon: Small constant added to denominator to avoid division by zero.
        reduction: Whether to compute "mean" or "sum" of normalized cross correlation
            of image pairs in batch. If "none", a 1-dimensional tensor is returned
            with length equal the batch size.

    Returns:
        Negative squared normalized cross correlation plus one.

    """

    if not isinstance(source, Tensor):
        raise TypeError("ncc_loss() 'source' must be tensor")
    if not isinstance(target, Tensor):
        raise TypeError("ncc_loss() 'target' must be tensor")
    if source.shape != target.shape:
        raise ValueError("ncc_loss() 'source' must have same shape as 'target'")

    source = source.reshape(source.shape[0], -1).float()
    target = target.reshape(source.shape[0], -1).float()
    print(f"{source.shape=},{target.shape=}")
    source_mean = source.mean(dim=1, keepdim=True)
    target_mean = target.mean(dim=1, keepdim=True)
    print(f"{source_mean.shape=},{target_mean.shape=}")
    x = source.sub(source_mean)
    y = target.sub(target_mean)
    print(f"{x.shape=},{y.shape=}")
    a = x.mul(y).sum(dim=1)
    b = x.square().sum(dim=1)
    c = y.square().sum(dim=1)
    print(f"{a.shape=},{b.shape=},{c.shape=}")

    loss = a.square_().div_(b.mul_(c).add_(epsilon)).neg_().add_(1)
    print(f"{loss.shape=},{mask.shape=}")
    loss = masked_loss(loss, mask, "ncc_loss")
    loss = reduce_loss(loss, reduction, mask)
    return loss

output:
source.shape=torch.Size([10, 1000]),target.shape=torch.Size([10, 1000])
source_mean.shape=torch.Size([10, 1]),target_mean.shape=torch.Size([10, 1])
x.shape=torch.Size([10, 1000]),y.shape=torch.Size([10, 1000])
a.shape=torch.Size([10]),b.shape=torch.Size([10]),c.shape=torch.Size([10])
loss.shape=torch.Size([10]),mask.shape=torch.Size([10, 10, 10, 10])

[aschuh-hf]

Thanks for catching this. On first glance, it looks like I may have forgotten the batch dimension when adding this function.