Repository: kornia/kornia

Issue:
losses.binary_focal_loss_with_logits is computed elementwise
### Describe the bug

Due to an apparent bug involving `unsqueeze`, batched `binary_focal_loss_with_logits` is computed pairwise along the batch dimension.

### Reproduction steps

```bash
I created a colab illustrating the issue. Please note how removing the two (un)squeezing lines seems to solve the issue. I can make a PR if wanted but I'm asking beforehand because I don't understand the rationale that got those unsqueezing operations to appear in the first place.

https://colab.research.google.com/drive/19OpQxpsJJfki7hV2Ht-3vDM6RO9c8x2W?usp=sharing
```


### Expected behavior

Loss for inputs of shape `(B, 1, ...)` should be of shape `(B, 1, ...)` or `(B, ...)` but not `(B, B, 1, ...)`.

### Environment

```shell
PyTorch version: 1.10.0.dev20210921+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.8.10 (default, Jun  4 2021, 15:09:15)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.11.0-1018-gcp-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-16GB
GPU 1: Tesla V100-SXM2-16GB
GPU 2: Tesla V100-SXM2-16GB
GPU 3: Tesla V100-SXM2-16GB
GPU 4: Tesla V100-SXM2-16GB
GPU 5: Tesla V100-SXM2-16GB
GPU 6: Tesla V100-SXM2-16GB
GPU 7: Tesla V100-SXM2-16GB

Nvidia driver version: 470.57.02
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.20.3
[pip3] pytorch-lightning==1.4.7
[pip3] torch==1.10.0.dev20210921+cu111
[pip3] torch-dimcheck==0.0.1
[pip3] torchmetrics==0.4.1
[pip3] torchvision==0.11.0.dev20210921+cu111
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.1.74              h6bb024c_0    nvidia
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.3.0           h06a4308_520  
[conda] mkl-service               2.4.0            py38h7f8727e_0  
[conda] mkl_fft                   1.3.0            py38h42c9631_2  
[conda] mkl_random                1.2.2            py38h51133e4_0  
[conda] mypy-extensions           0.4.3                    pypi_0    pypi
[conda] numpy                     1.19.5                   pypi_0    pypi
[conda] numpy-base                1.20.3           py38h74d4b33_0  
[conda] pytorch-lightning         1.4.7                    pypi_0    pypi
[conda] torch                     1.10.0.dev20210921+cu111          pypi_0    pypi
[conda] torchmetrics              0.4.1                    pypi_0    pypi
[conda] torchvision               0.11.0.dev20210921+cu111          pypi_0    pypi
```


### Additional context

_No response_


Discussion:
@jatentaki mind sending a PR with the fix. Besides, I''ll try to test tomorrow with a full end to end example for segmentation with the new training api :)

See: https://github.com/kornia/kornia/pull/1307

Resolving patch:
diff --git a/kornia/losses/focal.py b/kornia/losses/focal.py
--- a/kornia/losses/focal.py
+++ b/kornia/losses/focal.py
@@ -155,8 +155,8 @@ def binary_focal_loss_with_logits(
        - :math:`p_t` is the model's estimated probability for each class.
 
     Args:
-        input: input data tensor with shape :math:`(N, 1, *)`.
-        target: the target tensor with shape :math:`(N, 1, *)`.
+        input: input data tensor of arbitrary shape.
+        target: the target tensor with shape matching input.
         alpha: Weighting factor for the rare class :math:`\alpha \in [0, 1]`.
         gamma: Focusing parameter :math:`\gamma >= 0`.
         reduction: Specifies the reduction to apply to the
@@ -170,9 +170,8 @@ def binary_focal_loss_with_logits(
         the computed loss.
 
     Examples:
-        >>> num_classes = 1
         >>> kwargs = {"alpha": 0.25, "gamma": 2.0, "reduction": 'mean'}
-        >>> logits = torch.tensor([[[[6.325]]],[[[5.26]]],[[[87.49]]]])
+        >>> logits = torch.tensor([[[6.325]],[[5.26]],[[87.49]]])
         >>> labels = torch.tensor([[[1.]],[[1.]],[[0.]]])
         >>> binary_focal_loss_with_logits(logits, labels, **kwargs)
         tensor(4.6052)
@@ -188,13 +187,10 @@ def binary_focal_loss_with_logits(
         raise ValueError(f'Expected input batch_size ({input.size(0)}) to match target batch_size ({target.size(0)}).')
 
     probs = torch.sigmoid(input)
-    target = target.unsqueeze(dim=1)
     loss_tmp = -alpha * torch.pow((1.0 - probs + eps), gamma) * target * torch.log(probs + eps) - (
         1 - alpha
     ) * torch.pow(probs + eps, gamma) * (1.0 - target) * torch.log(1.0 - probs + eps)
 
-    loss_tmp = loss_tmp.squeeze(dim=1)
-
     if reduction == 'none':
         loss = loss_tmp
     elif reduction == 'mean':


Test patch:
diff --git a/test/losses/test_losses.py b/test/losses/test_losses.py
--- a/test/losses/test_losses.py
+++ b/test/losses/test_losses.py
@@ -11,8 +11,7 @@
 
 class TestBinaryFocalLossWithLogits:
     def test_smoke_none(self, device, dtype):
-        num_classes = 1
-        logits = torch.rand(2, num_classes, 3, 2, dtype=dtype, device=device)
+        logits = torch.rand(2, 3, 2, dtype=dtype, device=device)
         labels = torch.rand(2, 3, 2, dtype=dtype, device=device)
 
         assert kornia.losses.binary_focal_loss_with_logits(
@@ -20,8 +19,7 @@ def test_smoke_none(self, device, dtype):
         ).shape == (2, 3, 2)
 
     def test_smoke_sum(self, device, dtype):
-        num_classes = 1
-        logits = torch.rand(2, num_classes, 3, 2, dtype=dtype, device=device)
+        logits = torch.rand(2, 3, 2, dtype=dtype, device=device)
         labels = torch.rand(2, 3, 2, dtype=dtype, device=device)
 
         assert (
@@ -30,8 +28,7 @@ def test_smoke_sum(self, device, dtype):
         )
 
     def test_smoke_mean(self, device, dtype):
-        num_classes = 1
-        logits = torch.rand(2, num_classes, 3, 2, dtype=dtype, device=device)
+        logits = torch.rand(2, 3, 2, dtype=dtype, device=device)
         labels = torch.rand(2, 3, 2, dtype=dtype, device=device)
 
         assert (
@@ -40,8 +37,7 @@ def test_smoke_mean(self, device, dtype):
         )
 
     def test_smoke_mean_flat(self, device, dtype):
-        num_classes = 1
-        logits = torch.rand(2, num_classes, 3, 2, dtype=dtype, device=device)
+        logits = torch.rand(2, 3, 2, dtype=dtype, device=device)
         labels = torch.rand(2, 3, 2, dtype=dtype, device=device)
 
         assert (
@@ -50,8 +46,7 @@ def test_smoke_mean_flat(self, device, dtype):
         )
 
     def test_jit(self, device, dtype):
-        num_classes = 1
-        logits = torch.rand(2, num_classes, 3, 2, dtype=dtype, device=device)
+        logits = torch.rand(2, 3, 2, dtype=dtype, device=device)
         labels = torch.rand(2, 3, 2, dtype=dtype, device=device)
 
         op = kornia.losses.binary_focal_loss_with_logits
@@ -61,10 +56,9 @@ def test_jit(self, device, dtype):
         assert_close(actual, expected)
 
     def test_gradcheck(self, device):
-        num_classes = 1
         alpha, gamma = 0.5, 2.0  # for focal loss with logits
-        logits = torch.rand(2, num_classes, 3, 2).to(device)
-        labels = torch.rand(2, 1, 3, 2) * num_classes
+        logits = torch.rand(2, 3, 2).to(device)
+        labels = torch.rand(2, 1, 3, 2)
         labels = labels.to(device).long()
 
         logits = utils.tensor_to_gradcheck_var(logits)  # to var
@@ -73,8 +67,7 @@ def test_gradcheck(self, device):
         )
 
     def test_same_output(self, device, dtype):
-        num_classes = 1
-        logits = torch.rand(2, num_classes, 3, 2, dtype=dtype, device=device)
+        logits = torch.rand(2, 3, 2, dtype=dtype, device=device)
         labels = torch.rand(2, 3, 2, dtype=dtype, device=device)
 
         kwargs = {"alpha": 0.25, "gamma": 2.0, "reduction": 'mean'}