Repository: kornia/kornia

Issue:
Bug about the weight factor alpha of focal loss 
## 🐛 Bug

https://github.com/kornia/kornia/blob/cf8e854546110e3dac1caa35449429fc7af328e5/kornia/losses/focal.py#L86

Hi, I find someting wrong about the multi-class focal loss computation.

It seems like we are multiplying the loss of different classes by a same alpha factor. I think this is equal to just reduce lr, and cannot address the class imbalance issue. In the original paper, alpha is defined as another cased function just like class probability (Eq.2). 

Please take a look about it, thank you !



Discussion:
Up! I also think that this implementation is not correct original implementation in fvcore uses binary cross entropy and multiplies class 0 by alpha and class 1 by 1-alpha, so alpha is equivalent to weight parameter in CrossEntropy and I believe should also be normalized. (btw. in binary cross entropy you actually have this implemented right) 

EDIT: after some googling I decided to implement function on my own, it's equivalent to CrossEntropyLoss when gamma is 0, if somebody is interested go hear: https://github.com/Zhylkaaa/simpletransformers/blob/add_losses/simpletransformers/losses/focal_loss.py#L15
Yes. Alpha is a term that used as the weight in crossentropy and our implementation is probably wrong. @Zhylkaaa Are you willing to assist benchmarking against standard libraries like [this](https://github.com/qubvel/segmentation_models.pytorch/blob/a542c55575980aa19076bc53595c86df2b4103e3/segmentation_models_pytorch/losses/_functional.py#L36) then to throw a PR for a fix?
hey @shijianjian, my implementation is more about multi-class situation, so I don't know if there is a direct comparison. But I can try to help (but I am afraid it can take some time, it is a very busy month 😁).


Generally my implementation should be equivalent to qubvel, but you should expend single dimensional output to 2 dimensional vector (just [output, 1-output] for each prediction) , and set alpha to 1-alpha (in my implementation alpha is a weight for class 0)
> hey @shijianjian, my implementation is more about multi-class situation, so I don't know if there is a direct comparison. But I can try to help (but I am afraid it can take some time, it is a very busy month 😁).

Great. Thanks for your help.

Personally, I set the alpha value as per the class ratio of examples in a dataset. Basically, lower the number of examples, higher the weight as in ```nn.CrossEntropy(weight)```. Thus, I believe it might make sense to expose the ```alpha``` as both a tensor and a float. In such case, a float alpha would mean to handle binary classification while a tensor can be for both binary or multi-class. What do you think?

Another option can be the implementation of [torch-toolbelt](https://github.com/BloodAxe/pytorch-toolbelt/blob/920e03876805351ed5645e439a64074cb4f37589/pytorch_toolbelt/losses/functional.py#L82) which totally ignored the ```alpha``` for the multiclass cases.
Actually my implementation is exactly what you said. If you pass float to alpha it will be used as what I called one-vs-other mode (alpha is the weight for class 0, while 1-alpha is weight for the rest of classes) and if you pass a list (Iterable) than it will be used as normal weights (like in `nn.CrossEntropyLoss`)
This is implemented hear https://github.com/Zhylkaaa/simpletransformers/blob/add_losses/simpletransformers/losses/focal_loss.py#L100
I think weights can be beneficial (if I remember well original paper said that this is rather a good thing to have)
@jatentaki could you check if this is fixed in #1362 ?
It's not. My PR essentially replaces log(softmax(X)) with logsoftmax(x), a non-change if we used infinite precision arithmetic. I have not touched the core logic (equation) at all. I agree the alpha factor doesn't seem to be doing what it's supposed to.
@edgarriba I'd like to have a go at it, implementing a fix like what shijianjian and Zhylkaaa suggested.

What do you think?
Sure , please go ahead
Is it still the case?

Resolving patch:
diff --git a/kornia/losses/focal.py b/kornia/losses/focal.py
--- a/kornia/losses/focal.py
+++ b/kornia/losses/focal.py
@@ -1,11 +1,9 @@
 from __future__ import annotations
 
-import warnings
-
 import torch
 from torch import nn
 
-from kornia.core import Tensor
+from kornia.core import Tensor, tensor
 from kornia.core.check import KORNIA_CHECK, KORNIA_CHECK_IS_TENSOR, KORNIA_CHECK_SHAPE
 from kornia.utils.one_hot import one_hot
 
@@ -14,7 +12,12 @@
 
 
 def focal_loss(
-    input: Tensor, target: Tensor, alpha: float, gamma: float = 2.0, reduction: str = 'none', eps: float | None = None
+    input: Tensor,
+    target: Tensor,
+    alpha: float | None,
+    gamma: float = 2.0,
+    reduction: str = 'none',
+    weight: Tensor | None = None,
 ) -> Tensor:
     r"""Criterion that computes Focal loss.
 
@@ -29,7 +32,8 @@ def focal_loss(
 
     Args:
         input: logits tensor with shape :math:`(N, C, *)` where C = number of classes.
-        target: labels tensor with shape :math:`(N, *)` where each value is :math:`0 ≤ targets[i] ≤ C-1`.
+        target: labels tensor with shape :math:`(N, *)` where each value is an integer
+          representing correct classification :math:`target[i] \in [0, C)`.
         alpha: Weighting factor :math:`\alpha \in [0, 1]`.
         gamma: Focusing parameter :math:`\gamma >= 0`.
         reduction: Specifies the reduction to apply to the
@@ -37,49 +41,60 @@ def focal_loss(
           will be applied, ``'mean'``: the sum of the output will be divided by
           the number of elements in the output, ``'sum'``: the output will be
           summed.
-        eps: Deprecated: scalar to enforce numerical stabiliy. This is no longer used.
+        weight: weights for classes with shape :math:`(num_of_classes,)`.
 
     Return:
         the computed loss.
 
     Example:
-        >>> N = 5  # num_classes
-        >>> input = torch.randn(1, N, 3, 5, requires_grad=True)
-        >>> target = torch.empty(1, 3, 5, dtype=torch.long).random_(N)
-        >>> output = focal_loss(input, target, alpha=0.5, gamma=2.0, reduction='mean')
+        >>> C = 5  # num_classes
+        >>> input = torch.randn(1, C, 3, 5, requires_grad=True)
+        >>> target = torch.randint(C, (1, 3, 5))
+        >>> kwargs = {"alpha": 0.5, "gamma": 2.0, "reduction": 'mean'}
+        >>> output = focal_loss(input, target, **kwargs)
         >>> output.backward()
     """
-    if eps is not None and not torch.jit.is_scripting():
-        warnings.warn(
-            "`focal_loss` has been reworked for improved numerical stability "
-            "and the `eps` argument is no longer necessary",
-            DeprecationWarning,
-            stacklevel=2,
-        )
 
     KORNIA_CHECK_SHAPE(input, ["B", "C", "*"])
-
-    n = input.shape[0]
-    out_size = (n,) + input.shape[2:]
-
-    KORNIA_CHECK(target.shape[1:] == input.shape[2:], f'Expected target size {out_size}, got {target.size()}')
+    out_size = (input.shape[0],) + input.shape[2:]
+    KORNIA_CHECK(
+        (input.shape[0] == target.shape[0] and target.shape[1:] == input.shape[2:]),
+        f'Expected target size {out_size}, got {target.shape}',
+    )
     KORNIA_CHECK(
         input.device == target.device,
         f"input and target must be in the same device. Got: {input.device} and {target.device}",
     )
 
-    # compute softmax over the classes axis
-    input_soft: Tensor = input.softmax(1)
-    log_input_soft: Tensor = input.log_softmax(1)
-
     # create the labels one hot tensor
     target_one_hot: Tensor = one_hot(target, num_classes=input.shape[1], device=input.device, dtype=input.dtype)
 
+    # compute softmax over the classes axis
+    log_input_soft: Tensor = input.log_softmax(1)
+
     # compute the actual focal loss
-    weight = torch.pow(-input_soft + 1.0, gamma)
+    loss_tmp: Tensor = -torch.pow(1.0 - log_input_soft.exp(), gamma) * log_input_soft * target_one_hot
+
+    num_of_classes = input.shape[1]
+    boradcast_dims = [-1] + [1] * len(input.shape[2:])
+    if alpha is not None:
+        alpha_fac = tensor([1 - alpha] + [alpha] * (num_of_classes - 1), dtype=loss_tmp.dtype, device=loss_tmp.device)
+        alpha_fac = alpha_fac.view(boradcast_dims)
+        loss_tmp = alpha_fac * loss_tmp
+
+    if weight is not None:
+        KORNIA_CHECK_IS_TENSOR(weight, "weight must be Tensor or None.")
+        KORNIA_CHECK(
+            (weight.shape[0] == num_of_classes and weight.numel() == num_of_classes),
+            f'weight shape must be (num_of_classes,): ({num_of_classes},), got {weight.shape}',
+        )
+        KORNIA_CHECK(
+            weight.device == input.device,
+            f"weight and input must be in the same device. Got: {weight.device} and {input.device}",
+        )
 
-    focal = -alpha * weight * log_input_soft
-    loss_tmp = torch.einsum('bc...,bc...->b...', (target_one_hot, focal))
+        weight = weight.view(boradcast_dims)
+        loss_tmp = weight * loss_tmp
 
     if reduction == 'none':
         loss = loss_tmp
@@ -112,45 +127,48 @@ class FocalLoss(nn.Module):
           will be applied, ``'mean'``: the sum of the output will be divided by
           the number of elements in the output, ``'sum'``: the output will be
           summed.
-        eps: Deprecated: scalar to enforce numerical stability. This is no longer
-          used.
+        weight: weights for classes with shape :math:`(num_of_classes,)`.
 
     Shape:
         - Input: :math:`(N, C, *)` where C = number of classes.
-        - Target: :math:`(N, *)` where each value is
-          :math:`0 ≤ targets[i] ≤ C-1`.
+        - Target: :math:`(N, *)` where each value is an integer
+          representing correct classification :math:`target[i] \in [0, C)`.
 
     Example:
-        >>> N = 5  # num_classes
+        >>> C = 5  # num_classes
+        >>> input = torch.randn(1, C, 3, 5, requires_grad=True)
+        >>> target = torch.randint(C, (1, 3, 5))
         >>> kwargs = {"alpha": 0.5, "gamma": 2.0, "reduction": 'mean'}
         >>> criterion = FocalLoss(**kwargs)
-        >>> input = torch.randn(1, N, 3, 5, requires_grad=True)
-        >>> target = torch.empty(1, 3, 5, dtype=torch.long).random_(N)
         >>> output = criterion(input, target)
         >>> output.backward()
     """
 
-    def __init__(self, alpha: float, gamma: float = 2.0, reduction: str = 'none', eps: float | None = None) -> None:
+    def __init__(
+        self, alpha: float | None, gamma: float = 2.0, reduction: str = 'none', weight: Tensor | None = None
+    ) -> None:
         super().__init__()
-        self.alpha: float = alpha
+        self.alpha: float | None = alpha
         self.gamma: float = gamma
         self.reduction: str = reduction
-        self.eps: float | None = eps
+        self.weight: Tensor | None = weight
 
     def forward(self, input: Tensor, target: Tensor) -> Tensor:
-        return focal_loss(input, target, self.alpha, self.gamma, self.reduction, self.eps)
+        return focal_loss(input, target, self.alpha, self.gamma, self.reduction, self.weight)
 
 
 def binary_focal_loss_with_logits(
     input: Tensor,
     target: Tensor,
-    alpha: float = 0.25,
+    alpha: float | None = 0.25,
     gamma: float = 2.0,
     reduction: str = 'none',
-    eps: float | None = None,
     pos_weight: Tensor | None = None,
+    weight: Tensor | None = None,
 ) -> Tensor:
-    r"""Function that computes Binary Focal loss.
+    r"""Criterion that computes Binary Focal loss.
+
+    According to :cite:`lin2018focal`, the Focal loss is computed as follows:
 
     .. math::
 
@@ -160,61 +178,78 @@ def binary_focal_loss_with_logits(
        - :math:`p_t` is the model's estimated probability for each class.
 
     Args:
-        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]`.
+        input: logits tensor with shape :math:`(N, C, *)` where C = number of classes.
+        target: labels tensor with the same shape as input :math:`(N, C, *)`
+          where each value is between 0 and 1.
+        alpha: Weighting factor :math:`\alpha \in [0, 1]`.
         gamma: Focusing parameter :math:`\gamma >= 0`.
         reduction: Specifies the reduction to apply to the
           output: ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction
           will be applied, ``'mean'``: the sum of the output will be divided by
           the number of elements in the output, ``'sum'``: the output will be
           summed.
-        eps: Deprecated: scalar for numerically stability when dividing. This is no longer used.
-        pos_weight: a weight of positive examples.
-          It`s possible to trade off recall and precision by adding weights to positive examples.
-          Must be a vector with length equal to the number of classes.
+        pos_weight: a weight of positive examples with shape :math:`(num_of_classes,)`.
+          It is possible to trade off recall and precision by adding weights to positive examples.
+        weight: weights for classes with shape :math:`(num_of_classes,)`.
 
     Returns:
         the computed loss.
 
     Examples:
+        >>> C = 3  # num_classes
+        >>> input = torch.randn(1, C, 5, requires_grad=True)
+        >>> target = torch.randint(2, (1, C, 5))
         >>> kwargs = {"alpha": 0.25, "gamma": 2.0, "reduction": 'mean'}
-        >>> logits = torch.tensor([[[6.325]],[[5.26]],[[87.49]]])
-        >>> labels = torch.tensor([[[1.]],[[1.]],[[0.]]])
-        >>> binary_focal_loss_with_logits(logits, labels, **kwargs)
-        tensor(21.8725)
+        >>> output = binary_focal_loss_with_logits(input, target, **kwargs)
+        >>> output.backward()
     """
 
-    if eps is not None and not torch.jit.is_scripting():
-        warnings.warn(
-            "`binary_focal_loss_with_logits` has been reworked for improved numerical stability "
-            "and the `eps` argument is no longer necessary",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-
     KORNIA_CHECK_SHAPE(input, ["B", "C", "*"])
+    KORNIA_CHECK(input.shape == target.shape, f'Expected target size {input.shape}, got {target.shape}')
     KORNIA_CHECK(
-        input.shape[0] == target.shape[0],
-        f'Expected input batch_size ({input.shape[0]}) to match target batch_size ({target.shape[0]}).',
+        input.device == target.device,
+        f"input and target must be in the same device. Got: {input.device} and {target.device}",
     )
 
-    if pos_weight is None:
-        pos_weight = torch.ones([input.shape[1], *([1] * len(input.shape[2:]))], device=input.device, dtype=input.dtype)
+    log_probs_pos: Tensor = nn.functional.logsigmoid(input)
+    log_probs_neg: Tensor = nn.functional.logsigmoid(-input)
+
+    pos_term: Tensor = -log_probs_neg.exp().pow(gamma) * target * log_probs_pos
+    neg_term: Tensor = -log_probs_pos.exp().pow(gamma) * (1.0 - target) * log_probs_neg
+    if alpha is not None:
+        pos_term = alpha * pos_term
+        neg_term = (1.0 - alpha) * neg_term
+
+    num_of_classes = input.shape[1]
+    boradcast_dims = [-1] + [1] * len(input.shape[2:])
+    if pos_weight is not None:
+        KORNIA_CHECK_IS_TENSOR(pos_weight, "pos_weight must be Tensor or None.")
+        KORNIA_CHECK(
+            (pos_weight.shape[0] == num_of_classes and pos_weight.numel() == num_of_classes),
+            f'pos_weight shape must be (num_of_classes,): ({num_of_classes},), got {pos_weight.shape}',
+        )
+        KORNIA_CHECK(
+            pos_weight.device == input.device,
+            f"pos_weight and input must be in the same device. Got: {pos_weight.device} and {input.device}",
+        )
 
-    KORNIA_CHECK_IS_TENSOR(pos_weight)
-    KORNIA_CHECK(input.shape[1] == pos_weight.shape[0], "Expected pos_weight equals number of classes.")
+        pos_weight = pos_weight.view(boradcast_dims)
+        pos_term = pos_weight * pos_term
 
-    probs_pos = input.sigmoid()
-    probs_neg = (-input).sigmoid()
+    loss_tmp: Tensor = pos_term + neg_term
+    if weight is not None:
+        KORNIA_CHECK_IS_TENSOR(weight, "weight must be Tensor or None.")
+        KORNIA_CHECK(
+            (weight.shape[0] == num_of_classes and weight.numel() == num_of_classes),
+            f'weight shape must be (num_of_classes,): ({num_of_classes},), got {weight.shape}',
+        )
+        KORNIA_CHECK(
+            weight.device == input.device,
+            f"weight and input must be in the same device. Got: {weight.device} and {input.device}",
+        )
 
-    log_probs_pos = nn.functional.logsigmoid(input)
-    log_probs_neg = nn.functional.logsigmoid(-input)
-
-    loss_tmp = (
-        -alpha * pos_weight * probs_neg.pow(gamma) * target * log_probs_pos
-        - (1 - alpha) * probs_pos.pow(gamma) * (1.0 - target) * log_probs_neg
-    )
+        weight = weight.view(boradcast_dims)
+        loss_tmp = weight * loss_tmp
 
     if reduction == 'none':
         loss = loss_tmp
@@ -240,40 +275,48 @@ class BinaryFocalLossWithLogits(nn.Module):
        - :math:`p_t` is the model's estimated probability for each class.
 
     Args:
-        alpha: Weighting factor for the rare class :math:`\alpha \in [0, 1]`.
+        alpha: Weighting factor :math:`\alpha \in [0, 1]`.
         gamma: Focusing parameter :math:`\gamma >= 0`.
         reduction: Specifies the reduction to apply to the
           output: ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction
           will be applied, ``'mean'``: the sum of the output will be divided by
           the number of elements in the output, ``'sum'``: the output will be
           summed.
-        pos_weight: a weight of positive examples.
-          It`s possible to trade off recall and precision by adding weights to positive examples.
-          Must be a vector with length equal to the number of classes.
+        pos_weight: a weight of positive examples with shape :math:`(num_of_classes,)`.
+          It is possible to trade off recall and precision by adding weights to positive examples.
+        weight: weights for classes with shape :math:`(num_of_classes,)`.
 
     Shape:
-        - Input: :math:`(N, *)`.
-        - Target: :math:`(N, *)`.
+        - Input: :math:`(N, C, *)` where C = number of classes.
+        - Target: the same shape as Input :math:`(N, C, *)`
+          where each value is between 0 and 1.
 
     Examples:
+        >>> C = 3  # num_classes
+        >>> input = torch.randn(1, C, 5, requires_grad=True)
+        >>> target = torch.randint(2, (1, C, 5))
         >>> kwargs = {"alpha": 0.25, "gamma": 2.0, "reduction": 'mean'}
-        >>> loss = BinaryFocalLossWithLogits(**kwargs)
-        >>> input = torch.randn(1, 3, 5, requires_grad=True)
-        >>> target = torch.empty(1, 3, 5, dtype=torch.long).random_(2)
-        >>> output = loss(input, target)
+        >>> criterion = BinaryFocalLossWithLogits(**kwargs)
+        >>> output = criterion(input, target)
         >>> output.backward()
     """
 
     def __init__(
-        self, alpha: float, gamma: float = 2.0, reduction: str = 'none', pos_weight: Tensor | None = None
+        self,
+        alpha: float | None,
+        gamma: float = 2.0,
+        reduction: str = 'none',
+        pos_weight: Tensor | None = None,
+        weight: Tensor | None = None,
     ) -> None:
         super().__init__()
-        self.alpha: float = alpha
+        self.alpha: float | None = alpha
         self.gamma: float = gamma
         self.reduction: str = reduction
         self.pos_weight: Tensor | None = pos_weight
+        self.weight: Tensor | None = weight
 
     def forward(self, input: Tensor, target: Tensor) -> Tensor:
         return binary_focal_loss_with_logits(
-            input, target, self.alpha, self.gamma, self.reduction, pos_weight=self.pos_weight
+            input, target, self.alpha, self.gamma, self.reduction, self.pos_weight, self.weight
         )


Test patch:
diff --git a/test/integration/test_focal.py b/test/integration/test_focal.py
--- a/test/integration/test_focal.py
+++ b/test/integration/test_focal.py
@@ -18,7 +18,7 @@ class TestIntegrationFocalLoss:
     num_classes = 2
 
     # focal loss
-    alpha = 2.0
+    alpha = 0.5
     gamma = 2.0
 
     def generate_sample(self, base_target, std_val=0.1):
@@ -37,7 +37,11 @@ def test_conv2d_relu(self, device):
         for i in range(1, self.num_classes):
             target[..., i:-i, i:-i] = i
 
-        m = nn.Sequential(nn.Conv2d(1, self.num_classes, kernel_size=3, padding=1), nn.ReLU(True)).to(device)
+        m = nn.Sequential(
+            nn.Conv2d(1, self.num_classes // 2, kernel_size=3, padding=1),
+            nn.ReLU(True),
+            nn.Conv2d(self.num_classes // 2, self.num_classes, kernel_size=3, padding=1),
+        ).to(device)
         m.apply(self.init_weights)
 
         optimizer = optim.Adam(m.parameters(), lr=self.lr)
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
@@ -2,6 +2,7 @@
 
 import pytest
 import torch
+import torch.nn.functional as F
 from torch.autograd import gradcheck
 
 import kornia
@@ -10,59 +11,60 @@
 
 
 class TestBinaryFocalLossWithLogits:
-    def test_smoke_none(self, device, dtype):
+    @pytest.mark.parametrize('reduction', ['none', 'mean', 'sum'])
+    def test_value_same_as_torch_bce_loss(self, device, dtype, reduction):
         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(
-            logits, labels, alpha=0.5, gamma=2.0, reduction="none"
-        ).shape == (2, 3, 2)
-
-    def test_smoke_sum(self, device, dtype):
-        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(logits, labels, alpha=0.5, gamma=2.0, reduction="sum").shape
-            == ()
+        focal_equivalent_bce_loss = kornia.losses.binary_focal_loss_with_logits(
+            logits, labels, alpha=None, gamma=0, reduction=reduction
         )
+        torch_bce_loss = F.binary_cross_entropy_with_logits(logits, labels, reduction=reduction)
+        assert_close(focal_equivalent_bce_loss, torch_bce_loss)
 
-    def test_smoke_mean(self, device, dtype):
-        logits = torch.rand(2, 3, 2, dtype=dtype, device=device)
-        labels = torch.rand(2, 3, 2, dtype=dtype, device=device)
+    @pytest.mark.parametrize('reduction', ['none', 'mean', 'sum'])
+    def test_value_same_as_torch_bce_loss_pos_weight_weight(self, device, dtype, reduction):
+        num_classes = 3
+        logits = torch.rand(2, num_classes, 2, dtype=dtype, device=device)
+        labels = torch.rand(2, num_classes, 2, dtype=dtype, device=device)
+
+        pos_weight = torch.rand(num_classes, 1, dtype=dtype, device=device)
+        weight = torch.rand(num_classes, 1, dtype=dtype, device=device)
 
-        assert (
-            kornia.losses.binary_focal_loss_with_logits(logits, labels, alpha=0.5, gamma=2.0, reduction="mean").shape
-            == ()
+        focal_equivalent_bce_loss = kornia.losses.binary_focal_loss_with_logits(
+            logits, labels, alpha=None, gamma=0, reduction=reduction, pos_weight=pos_weight, weight=weight
         )
+        torch_bce_loss = F.binary_cross_entropy_with_logits(
+            logits, labels, reduction=reduction, pos_weight=pos_weight, weight=weight
+        )
+        assert_close(focal_equivalent_bce_loss, torch_bce_loss)
 
-    def test_smoke_mean_flat(self, device, dtype):
+    @pytest.mark.parametrize('reduction,expected_shape', [('none', (2, 3, 2)), ('mean', ()), ('sum', ())])
+    @pytest.mark.parametrize('alpha', [None, 0.2, 0.5])
+    @pytest.mark.parametrize('gamma', [0.0, 1.0, 2.0])
+    def test_shape_alpha_gamma(self, device, dtype, reduction, expected_shape, alpha, gamma):
         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(logits, labels, alpha=0.5, gamma=2.0, reduction="mean").shape
-            == ()
-        )
+        actual_shape = kornia.losses.binary_focal_loss_with_logits(
+            logits, labels, alpha=alpha, gamma=gamma, reduction=reduction
+        ).shape
+        assert actual_shape == expected_shape
 
-    @pytest.mark.parametrize("pos_weight", [None, 1, 5])
-    def test_smoke_pos_weight(self, device, dtype, pos_weight):
+    @pytest.mark.parametrize('reduction,expected_shape', [('none', (2, 3, 2)), ('mean', ()), ('sum', ())])
+    @pytest.mark.parametrize("pos_weight", [None, (1, 2, 5)])
+    @pytest.mark.parametrize("weight", [None, (0.2, 0.5, 0.8)])
+    def test_shape_pos_weight_weight(self, device, dtype, reduction, expected_shape, pos_weight, weight):
         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(
-                logits,
-                labels,
-                alpha=0.5,
-                gamma=2.0,
-                reduction="mean",
-                pos_weight=None
-                if pos_weight is None
-                else torch.full([logits.size(1), 1], pos_weight, dtype=dtype, device=device),
-            ).shape
-            == ()
-        )
+        pos_weight = None if pos_weight is None else torch.tensor(pos_weight, dtype=dtype, device=device)
+        weight = None if weight is None else torch.tensor(weight, dtype=dtype, device=device)
+
+        actual_shape = kornia.losses.binary_focal_loss_with_logits(
+            logits, labels, alpha=0.8, gamma=0.5, reduction=reduction, pos_weight=pos_weight, weight=weight
+        ).shape
+        assert actual_shape == expected_shape
 
     def test_dynamo(self, device, dtype, torch_optimizer):
         logits = torch.rand(2, 3, 2, dtype=dtype, device=device)
@@ -70,113 +72,106 @@ def test_dynamo(self, device, dtype, torch_optimizer):
 
         op = kornia.losses.binary_focal_loss_with_logits
         op_optimized = torch_optimizer(op)
-        actual = op_optimized(logits, labels, alpha=0.5, gamma=2.0, reduction="none")
-        expected = op(logits, labels, alpha=0.5, gamma=2.0, reduction="none")
+
+        args = (0.25, 2.0)
+        actual = op_optimized(logits, labels, *args)
+        expected = op(logits, labels, *args)
         assert_close(actual, expected)
 
     def test_gradcheck(self, device, dtype):
-        alpha, gamma = 0.5, 2.0  # for focal loss with logits
         logits = torch.rand(2, 3, 2).to(device, dtype)
-        labels = torch.rand(2, 1, 3, 2)
-        labels = labels.to(device, dtype).long()
-
         logits = tensor_to_gradcheck_var(logits)  # to var
-        assert gradcheck(
-            kornia.losses.binary_focal_loss_with_logits,
-            (logits, labels, alpha, gamma),
-            raise_exception=True,
-            fast_mode=True,
-        )
+        labels = torch.rand(2, 3, 2).to(device, dtype)
+
+        args = (0.25, 2.0)
+        op = kornia.losses.binary_focal_loss_with_logits
+        assert gradcheck(op, (logits, labels, *args), raise_exception=True, fast_mode=True)
 
-    def test_same_output(self, device, dtype):
+    def test_module(self, device, dtype):
         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"}
-
-        assert kornia.losses.binary_focal_loss_with_logits(
-            logits, labels, **kwargs
-        ) == kornia.losses.BinaryFocalLossWithLogits(**kwargs)(logits, labels)
+        args = (0.25, 2.0)
+        op = kornia.losses.binary_focal_loss_with_logits
+        op_module = kornia.losses.BinaryFocalLossWithLogits(*args)
+        assert_close(op_module(logits, labels), op(logits, labels, *args))
 
     def test_numeric_stability(self, device, dtype):
         logits = torch.tensor([[100.0, -100]], dtype=dtype, device=device)
         labels = torch.tensor([[1.0, 0.0]], dtype=dtype, device=device)
-        expected = torch.tensor([[0.0, 0.0]], dtype=dtype, device=device)
 
-        kwargs = {"alpha": 0.25, "gamma": 2.0, "reduction": 'none'}
-        actual = kornia.losses.binary_focal_loss_with_logits(logits, labels, **kwargs)
-        assert_close(actual, expected, atol=1e-3, rtol=1e-3)
+        args = (0.25, 2.0)
+        actual = kornia.losses.binary_focal_loss_with_logits(logits, labels, *args)
+        expected = torch.tensor([[0.0, 0.0]], dtype=dtype, device=device)
+        assert_close(actual, expected)
 
 
 class TestFocalLoss:
-    def test_smoke_none(self, device, dtype):
+    @pytest.mark.parametrize('reduction,expected_shape', [('none', (2, 3, 3, 2)), ('mean', ()), ('sum', ())])
+    @pytest.mark.parametrize('alpha', [None, 0.2, 0.5])
+    @pytest.mark.parametrize('gamma', [0.0, 1.0, 2.0])
+    def test_shape_alpha_gamma(self, device, dtype, reduction, expected_shape, alpha, gamma):
         num_classes = 3
         logits = torch.rand(2, num_classes, 3, 2, device=device, dtype=dtype)
-        labels = torch.rand(2, 3, 2) * num_classes
-        labels = labels.to(device).long()
+        labels = torch.randint(num_classes, (2, 3, 2), device=device)
 
-        assert kornia.losses.focal_loss(logits, labels, alpha=0.5, gamma=2.0, reduction="none").shape == (2, 3, 2)
+        actual_shape = kornia.losses.focal_loss(logits, labels, alpha=alpha, gamma=gamma, reduction=reduction).shape
+        assert actual_shape == expected_shape
 
-    def test_smoke_sum(self, device, dtype):
+    @pytest.mark.parametrize('reduction,expected_shape', [('none', (2, 3)), ('mean', ()), ('sum', ())])
+    def test_shape_target_with_only_one_dim(self, device, dtype, reduction, expected_shape):
         num_classes = 3
-        logits = torch.rand(2, num_classes, 3, 2, device=device, dtype=dtype)
-        labels = torch.rand(2, 3, 2) * num_classes
-        labels = labels.to(device).long()
+        logits = torch.rand(2, num_classes, device=device, dtype=dtype)
+        labels = torch.randint(num_classes, (2,), device=device)
 
-        assert kornia.losses.focal_loss(logits, labels, alpha=0.5, gamma=2.0, reduction="sum").shape == ()
+        actual_shape = kornia.losses.focal_loss(logits, labels, alpha=0.1, gamma=1.5, reduction=reduction).shape
+        assert actual_shape == expected_shape
 
-    def test_smoke_mean(self, device, dtype):
+    @pytest.mark.parametrize('reduction,expected_shape', [('none', (2, 3, 3, 2)), ('mean', ()), ('sum', ())])
+    @pytest.mark.parametrize("weight", [None, (0.2, 0.5, 0.8)])
+    def test_shape_weight(self, device, dtype, reduction, expected_shape, weight):
         num_classes = 3
         logits = torch.rand(2, num_classes, 3, 2, device=device, dtype=dtype)
-        labels = torch.rand(2, 3, 2) * num_classes
-        labels = labels.to(device).long()
-
-        assert kornia.losses.focal_loss(logits, labels, alpha=0.5, gamma=2.0, reduction="mean").shape == ()
-
-    def test_smoke_mean_flat(self, device, dtype):
-        num_classes = 3
-        logits = torch.rand(2, num_classes, device=device, dtype=dtype)
-        labels = torch.rand(2) * num_classes
-        labels = labels.to(device).long()
-        assert kornia.losses.focal_loss(logits, labels, alpha=0.5, gamma=2.0, reduction="mean").shape == ()
+        labels = torch.randint(num_classes, (2, 3, 2), device=device)
 
-    def test_gradcheck(self, device, dtype):
-        num_classes = 3
-        alpha, gamma = 0.5, 2.0  # for focal loss
-        logits = torch.rand(2, num_classes, 3, 2, device=device, dtype=dtype)
-        labels = torch.rand(2, 3, 2) * num_classes
-        labels = labels.to(device).long()
+        weight = None if weight is None else torch.tensor(weight, dtype=dtype, device=device)
 
-        logits = tensor_to_gradcheck_var(logits)  # to var
-        assert gradcheck(kornia.losses.focal_loss, (logits, labels, alpha, gamma), raise_exception=True, fast_mode=True)
+        actual_shape = kornia.losses.focal_loss(
+            logits, labels, alpha=0.8, gamma=0.5, reduction=reduction, weight=weight
+        ).shape
+        assert actual_shape == expected_shape
 
     def test_dynamo(self, device, dtype, torch_optimizer):
         num_classes = 3
-        params = (0.5, 2.0)
         logits = torch.rand(2, num_classes, device=device, dtype=dtype)
-        labels = torch.rand(2) * num_classes
-        labels = labels.to(device).long()
+        labels = torch.randint(num_classes, (2,), device=device)
 
         op = kornia.losses.focal_loss
         op_optimized = torch_optimizer(op)
 
-        actual = op_optimized(logits, labels, *params)
-        expected = op(logits, labels, *params)
+        args = (0.25, 2.0)
+        actual = op_optimized(logits, labels, *args)
+        expected = op(logits, labels, *args)
         assert_close(actual, expected)
 
+    def test_gradcheck(self, device, dtype):
+        num_classes = 3
+        logits = torch.rand(2, num_classes, 3, 2, device=device, dtype=dtype)
+        logits = tensor_to_gradcheck_var(logits)  # to var
+        labels = torch.randint(num_classes, (2, 3, 2), device=device)
+
+        args = (0.25, 2.0)
+        assert gradcheck(kornia.losses.focal_loss, (logits, labels, *args), raise_exception=True, fast_mode=True)
+
     def test_module(self, device, dtype):
         num_classes = 3
-        params = (0.5, 2.0)
         logits = torch.rand(2, num_classes, device=device, dtype=dtype)
-        labels = torch.rand(2) * num_classes
-        labels = labels.to(device).long()
+        labels = torch.randint(num_classes, (2,), device=device)
 
+        args = (0.25, 2.0)
         op = kornia.losses.focal_loss
-        op_module = kornia.losses.FocalLoss(*params)
-
-        actual = op_module(logits, labels)
-        expected = op(logits, labels, *params)
-        assert_close(actual, expected)
+        op_module = kornia.losses.FocalLoss(*args)
+        assert_close(op_module(logits, labels), op(logits, labels, *args))
 
 
 class TestTverskyLoss: