Focal loss with reduction=none returns an incorrect shape #3185
Description
Describe the bug
Note: I've already submitted a patch for this in #3041
The MONAI implementation of FocalLoss(reduction='none')(input, target) does not return a Tensor with the same shape as input. This deviates from the behavior of torch.nn.BCEWithLogitsLoss and other similar loss functions.
Currently using FocalLoss with reduction='none' will accept an input and target tensor with shape (B, N, *DIMS) but then it returns a tensor of shape (B, N). This is inconsistent with BCEWithLogitsLoss, which this form of FocalLoss is an extension of (There is another variant where it extents Categorical Cross Entropy, but that does not seem to be implemented here), which would return a tensor with shape (B, N, *DIMS) when reduction='none'.
This can be seen with the following code:
>>> import torch
>>> from monai.losses import FocalLoss
>>> from torch.nn import BCEWithLogitsLoss
>>> shape = B, N, *DIMS = 2, 3, 5, 7, 11
>>> input = torch.rand(*shape)
>>> target = torch.rand(*shape)
>>> # Demonstrate equivalence to BCE when gamma=0
>>> fl_g0_criterion = FocalLoss(reduction='none', gamma=0)
>>> bce_criterion = BCEWithLogitsLoss(reduction='none')
>>> fl_g0_loss = fl_g0_criterion(input, target)
>>> bce_loss = bce_criterion(input, target)
>>> print('bce_loss.shape = {!r}'.format(bce_loss.shape))
>>> print('fl_g0_loss.shape = {!r}'.format(fl_g0_loss.shape))The current code will produce:
bce_loss.shape = torch.Size([2, 3, 5, 7, 11])
fl_g0_loss.shape = torch.Size([2, 3])
Expected behavior
The expected shapes between bce and focal loss should match
bce_loss.shape = torch.Size([2, 3, 5, 7, 11])
fl_g0_loss.shape = torch.Size([2, 3, 5, 7, 11])
Environment
Ensuring you use the relevant python executable, please paste the output of:
(pyenv3.8.6) jon.crall@yardrat:~/code/MONAI$ python -c 'import monai; monai.config.print_debug_info()'
================================
Printing MONAI config...
================================
MONAI version: 0.7.0+83.g406651a5
Numpy version: 1.21.2
Pytorch version: 1.10.0+cu113
MONAI flags: HAS_EXT = False, USE_COMPILED = False
MONAI rev id: 406651a5825635b75f5669cdf75aa8de90479fc0
Optional dependencies:
Pytorch Ignite version: 0.4.5
Nibabel version: 3.2.1
scikit-image version: 0.18.2
Pillow version: 8.3.1
Tensorboard version: 2.6.0
gdown version: 3.13.1
TorchVision version: 0.11.1+cu113
tqdm version: 4.62.0
lmdb version: 1.2.1
psutil version: 5.8.0
pandas version: 1.3.1
einops version: 0.3.0
transformers version: 4.11.0
mlflow version: NOT INSTALLED or UNKNOWN VERSION.
For details about installing the optional dependencies, please visit:
https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies
================================
Printing system config...
================================
System: Linux
Linux version: Ubuntu 20.04.3 LTS
Platform: Linux-5.11.0-27-generic-x86_64-with-glibc2.2.5
Processor: x86_64
Machine: x86_64
Python version: 3.8.6
Process name: python
Command: ['python', '-c', 'import monai; monai.config.print_debug_info()']
Open files: []
Num physical CPUs: 8
Num logical CPUs: 16
Num usable CPUs: 16
CPU usage (%): [14.4, 5.8, 5.8, 5.8, 59.7, 5.0, 4.3, 5.0, 6.4, 89.2, 8.6, 5.8, 15.7, 35.5, 6.5, 5.7]
CPU freq. (MHz): 1578
Load avg. in last 1, 5, 15 mins (%): [9.2, 8.9, 7.2]
Disk usage (%): 56.5
Avg. sensor temp. (Celsius): UNKNOWN for given OS
Total physical memory (GB): 62.6
Available memory (GB): 54.5
Used memory (GB): 7.2
================================
Printing GPU config...
================================
Num GPUs: 1
Has CUDA: True
CUDA version: 11.3
cuDNN enabled: True
cuDNN version: 8200
Current device: 0
Library compiled for CUDA architectures: ['sm_37', 'sm_50', 'sm_60', 'sm_70', 'sm_75', 'sm_80', 'sm_86']
GPU 0 Name: Quadro RTX 5000
GPU 0 Is integrated: False
GPU 0 Is multi GPU board: False
GPU 0 Multi processor count: 48
GPU 0 Total memory (GB): 15.7
GPU 0 CUDA capability (maj.min): 7.5