Focal Loss returning 0.0 with include_background=False (7/July)

[SteEcker]

Describe the bug
Focal Loss returns 0 for all entries when input consists of only one channel. Tried running the spleen 3D segmentation tutorial with DiceFocalLoss(include_background=False). Printing out dice and focal loss separately shows the problem.

I think the issue is that logpt = F.log_softmax(i, dim=1) returns 0 for arbritrary inputs. Should this be F.logsigmoid() in this case?

To Reproduce

input = torch.randn(1, 1, 10)
F.log_softmax(input, dim=1)

tensor([[[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]])

vs

input = torch.randn(1, 2, 10)
F.log_softmax(input, dim=1)

input = torch.randn(1, 2, 10)...
tensor([[[-2.6699, -0.4457, -1.1688, -0.3480, -0.9432, -0.8112, -0.2835,
          -2.2028, -1.3333, -1.8177],
         [-0.0718, -1.0226, -0.3722, -1.2245, -0.4933, -0.5876, -1.3990,
          -0.1171, -0.3060, -0.1772]]])

Expected behavior
Focal Loss should return correct log.probs. for single channel inputs.

Environment

Printing MONAI config...

MONAI version: 0.6.dev2126
Numpy version: 1.19.5
Pytorch version: 1.9.0+cu102
MONAI flags: HAS_EXT = False, USE_COMPILED = False
MONAI rev id: 2ad54662de25e9a964c33327f7f2f178655573ef

Optional dependencies:
Pytorch Ignite version: NOT INSTALLED or UNKNOWN VERSION.
Nibabel version: 3.0.2
scikit-image version: 0.16.2
Pillow version: 7.1.2
Tensorboard version: 2.4.1
gdown version: 3.6.4
TorchVision version: 0.10.0+cu102
ITK version: NOT INSTALLED or UNKNOWN VERSION.
tqdm version: 4.41.1
lmdb version: 0.99
psutil version: 5.4.8
pandas version: 1.1.5

[Nic-Ma]

Hi @SteEcker ,

Thanks for your issue report!

@yiheng-wang-nv , could you please help check this issue?

Thanks in advance.

[yiheng-wang-nv]

Hi @SteEcker , could you please clarify what is the number of classes in your case?
As the docsting shown in: https://github.com/Project-MONAI/MONAI/blob/dev/monai/losses/focal_loss.py#L79
and https://github.com/Project-MONAI/MONAI/blob/dev/monai/losses/focal_loss.py#L42
The channel with index 0 represents the background, and if the input only have two classes, the channel with index 1 represents the only foreground, thus why set include_background=False?

[SteEcker]

Hi @yiheng-wang-nv and @Nic-Ma

Yes I am referring to a case where there are only two classes: background and foreground.

why set include_background=False?

In my understanding, even in a two class case, if the non-background segmentations are small compared to the total image size, the DICE loss can underestimate the actual loss. Maybe I misunderstood something. Here is a toy example to illustrate my question.

As you said input and target shape should be BNH[WD]. Assuming a 5x5 image with B = 1 my input and target tensors are of size (1,1,5,5) .

import torch 
from monai.losses import DiceLoss

background = torch.ones(1,1,5,5)
foreground = torch.zeros(1,1,5,5)

target = torch.cat((background,foreground), dim=1 )
input = torch.cat((background, foreground), dim=1)
target[:,0, 2,2] = 0 
target[:,1, 2,2] = 1 

print(
    f""" 
    TARGET: 
    {target}
    
    INPUT: 
    {input}
     """
)
dice_loss = DiceLoss(include_background=True)
loss = dice_loss(input, target)
print(f"DICE Loss with background: {loss}")

dice_loss = DiceLoss(include_background=False)
loss = dice_loss(input, target)
print(f"DICE Loss without background: {loss}")

Results in

 TARGET: 
    tensor([[[[1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1.],
          [1., 1., 0., 1., 1.],
          [1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1.]],

         [[0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0.],
          [0., 0., 1., 0., 0.],
          [0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0.]]]])
    
    INPUT: 
    tensor([[[[1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1.]],

         [[0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0.]]]])
     
DICE Loss with background: 0.5101990699768066
DICE Loss without background: 0.9999899864196777

As you can see it works as intended with DiceLoss. However imo it breaks down when calculating the Focal Loss as a logsoftmax across a single channel will always return 0. I hope this cleared things up. Thanks for the help and let me know if I am missing something obvious here 😄

[wyli]

I think you're right @SteEcker, the original paper uses a sigmoid https://arxiv.org/pdf/1708.02002.pdf, I have a better implementation, will submit a fix soon