Repository: kornia/kornia

Issue:
Improve binary focal loss to trade off recall and precision by adding weights to positive examples like torch.nn.BCEWithLogitsLoss
## 🚀 Feature
Improve binary focal loss to trade off recall and precision by adding weights to positive examples. Just like `pos_weight` parameter in torch.nn.BCEWithLogitsLoss. `pos_weight` must be a vector with length equal to the number of classes.



Discussion:
@xiaoyuan0203 Open to send a PR ?
@xiaoyuan0203 @edgarriba Looks interesting! Can I take it?
@MrShevan sure ! go ahead
> @xiaoyuan0203 Open to send a PR ?

Sorry for reply to you until now. Could you tell me how to do it? 
Oh sorry. I see you are interested in creating your own PR. Let me know if it's free


It's free of cource. It's my first time to take part in open-source,so I just want to know how to do. If you are interested in make a pull request, you can do it. Do you mean you want to improve the focal loss?
@xiaoyuan0203 i have other suggestions for losses to add e.g for semantic segmentation. or the ones from this model we are going to include: https://github.com/xavysp/LDC
@xiaoyuan0203 Well, I can help you get started :)
We can try to make this PR together. You can fork this and create your own branch. [Here](https://devblog.pytorchlightning.ai/quick-contribution-guide-86d977171b3a) it is well written in paragraph `Prepare your Workspace`. And then I'll join you)
you can also check our less graphical contributor guides: https://github.com/kornia/kornia/blob/master/CONTRIBUTING.rst#developing-kornia
@MrShevan Here? https://github.com/xiaoyuan0203/kornia/tree/feat/focal_loss
Yep, great)
I have done the following steps:
1. I cloned your forked repo into my local and switched into your branch. 
I use PyCharm as IDE, so it was simple to do the initial setup
<img width="1440" alt="Снимок экрана 2022-06-30 в 10 15 15" src="https://user-images.githubusercontent.com/32237302/176621692-0596d00a-e3dc-41fb-9343-a4f48e11e295.png">

and switch to `feat/focal_loss`
<img width="373" alt="Снимок экрана 2022-06-30 в 10 33 33" src="https://user-images.githubusercontent.com/32237302/176621918-9ccb9a2f-7be2-4df4-80d3-1c2d9024e720.png">

2. The next step is setting your python interpreter and run all tests **before starting coding**. :)

I ran all commands from [this guide](https://github.com/kornia/kornia/blob/master/CONTRIBUTING.rst#developing-kornia) in 4:
```
source ./path.bash.inc
python setup.py develop
python -c "import kornia; print(kornia.__version__)"
```
And then install all requirements to install pytest and etc.:
```
pip install --ignore-installed -r requirements/dev.txt
```

Now I can run all related to our PR tests (cause it is important to make sure that all the tests that we will change worked initially):
<img width="1440" alt="Снимок экрана 2022-06-30 в 10 19 05" src="https://user-images.githubusercontent.com/32237302/176623161-bb2e9687-a9ee-44ab-855c-3d44866fa321.png">

And they all passed :)
<img width="1418" alt="Снимок экрана 2022-06-30 в 10 52 32" src="https://user-images.githubusercontent.com/32237302/176623504-5107f2e5-b0bf-4cc2-a7ec-670878cedb3b.png">

Well, now I'm ready for coding 

@MrShevan Thanks a lot! Is it a good habit to use my usual conda environment instead of environment from the `kornia/setup.py`? How about other project I take part in?
Well, I think If the tests run, it doesn't matter. But sometimes tests too complicated and require a lot of installed python libraries, in such case easy create new interpreter as I did or updateConda (command from [this guide](https://github.com/kornia/kornia/blob/master/CONTRIBUTING.rst#developing-kornia) in 4: `./setup_dev_env.sh`)

Do you mind if I open the first PR from your branch with my draft changes?)
It doesn't matter. Just go ahead.

Resolving patch:
diff --git a/kornia/losses/focal.py b/kornia/losses/focal.py
--- a/kornia/losses/focal.py
+++ b/kornia/losses/focal.py
@@ -3,8 +3,9 @@
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
+from kornia.core import Tensor
+from kornia.testing import KORNIA_CHECK, KORNIA_CHECK_IS_TENSOR, KORNIA_CHECK_SHAPE
 from kornia.utils.one_hot import one_hot
 
 # based on:
@@ -12,13 +13,13 @@
 
 
 def focal_loss(
-    input: torch.Tensor,
-    target: torch.Tensor,
+    input: Tensor,
+    target: Tensor,
     alpha: float,
     gamma: float = 2.0,
     reduction: str = 'none',
     eps: Optional[float] = None,
-) -> torch.Tensor:
+) -> Tensor:
     r"""Criterion that computes Focal loss.
 
     According to :cite:`lin2018focal`, the Focal loss is computed as follows:
@@ -60,29 +61,23 @@ def focal_loss(
             stacklevel=2,
         )
 
-    if not isinstance(input, torch.Tensor):
-        raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")
+    KORNIA_CHECK_SHAPE(input, ["B", "C", "*"])
 
-    if not len(input.shape) >= 2:
-        raise ValueError(f"Invalid input shape, we expect BxCx*. Got: {input.shape}")
+    n = input.shape[0]
+    out_size = (n,) + input.shape[2:]
 
-    if input.size(0) != target.size(0):
-        raise ValueError(f'Expected input batch_size ({input.size(0)}) to match target batch_size ({target.size(0)}).')
-
-    n = input.size(0)
-    out_size = (n,) + input.size()[2:]
-    if target.size()[1:] != input.size()[2:]:
-        raise ValueError(f'Expected target size {out_size}, got {target.size()}')
-
-    if not input.device == target.device:
-        raise ValueError(f"input and target must be in the same device. Got: {input.device} and {target.device}")
+    KORNIA_CHECK(target.shape[1:] == input.shape[2:], f'Expected target size {out_size}, got {target.size()}')
+    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: torch.Tensor = F.softmax(input, dim=1)
-    log_input_soft: torch.Tensor = F.log_softmax(input, dim=1)
+    input_soft: Tensor = input.softmax(1)
+    log_input_soft: Tensor = input.log_softmax(1)
 
     # create the labels one hot tensor
-    target_one_hot: torch.Tensor = one_hot(target, num_classes=input.shape[1], device=input.device, dtype=input.dtype)
+    target_one_hot: Tensor = one_hot(target, num_classes=input.shape[1], device=input.device, dtype=input.dtype)
 
     # compute the actual focal loss
     weight = torch.pow(-input_soft + 1.0, gamma)
@@ -146,18 +141,19 @@ def __init__(self, alpha: float, gamma: float = 2.0, reduction: str = 'none', ep
         self.reduction: str = reduction
         self.eps: Optional[float] = eps
 
-    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
         return focal_loss(input, target, self.alpha, self.gamma, self.reduction, self.eps)
 
 
 def binary_focal_loss_with_logits(
-    input: torch.Tensor,
-    target: torch.Tensor,
+    input: Tensor,
+    target: Tensor,
     alpha: float = 0.25,
     gamma: float = 2.0,
     reduction: str = 'none',
     eps: Optional[float] = None,
-) -> torch.Tensor:
+    pos_weight: Optional[Tensor] = None,
+) -> Tensor:
     r"""Function that computes Binary Focal loss.
 
     .. math::
@@ -178,6 +174,9 @@ def binary_focal_loss_with_logits(
           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.
 
     Returns:
         the computed loss.
@@ -198,20 +197,25 @@ def binary_focal_loss_with_logits(
             stacklevel=2,
         )
 
-    if not isinstance(input, torch.Tensor):
-        raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")
+    KORNIA_CHECK_SHAPE(input, ["B", "C", "*"])
+    KORNIA_CHECK(
+        input.shape[0] == target.shape[0],
+        f'Expected input batch_size ({input.shape[0]}) to match target batch_size ({target.shape[0]}).',
+    )
 
-    if not len(input.shape) >= 2:
-        raise ValueError(f"Invalid input shape, we expect BxCx*. Got: {input.shape}")
+    if pos_weight is None:
+        pos_weight = torch.ones(input.shape[-1], device=input.device, dtype=input.dtype)
 
-    if input.size(0) != target.size(0):
-        raise ValueError(f'Expected input batch_size ({input.size(0)}) to match target batch_size ({target.size(0)}).')
+    KORNIA_CHECK_IS_TENSOR(pos_weight)
+    KORNIA_CHECK(input.shape[-1] == pos_weight.shape[0], "Expected pos_weight equals number of classes.")
 
-    probs_pos = torch.sigmoid(input)
+    probs_pos = input.sigmoid()
     probs_neg = torch.sigmoid(-input)
-    loss_tmp = -alpha * torch.pow(probs_neg, gamma) * target * F.logsigmoid(input) - (1 - alpha) * torch.pow(
-        probs_pos, gamma
-    ) * (1.0 - target) * F.logsigmoid(-input)
+
+    loss_tmp = (
+        -alpha * pos_weight * probs_neg.pow(gamma) * target * input.sigmoid().log()
+        - (1 - alpha) * torch.pow(probs_pos, gamma) * (1.0 - target) * (-input).sigmoid().log()
+    )
 
     if reduction == 'none':
         loss = loss_tmp
@@ -237,13 +241,16 @@ 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 for the rare class :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.
 
     Shape:
         - Input: :math:`(N, *)`.
@@ -258,11 +265,16 @@ class BinaryFocalLossWithLogits(nn.Module):
         >>> output.backward()
     """
 
-    def __init__(self, alpha: float, gamma: float = 2.0, reduction: str = 'none') -> None:
+    def __init__(
+        self, alpha: float, gamma: float = 2.0, reduction: str = 'none', pos_weight: Optional[Tensor] = None
+    ) -> None:
         super().__init__()
         self.alpha: float = alpha
         self.gamma: float = gamma
         self.reduction: str = reduction
+        self.pos_weight: Optional[Tensor] = pos_weight
 
-    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        return binary_focal_loss_with_logits(input, target, self.alpha, self.gamma, self.reduction)
+    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
+        )


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
@@ -45,6 +45,25 @@ def test_smoke_mean_flat(self, device, dtype):
             == ()
         )
 
+    @pytest.mark.parametrize("pos_weight", [None, 1, 5])
+    def test_smoke_pos_weight(self, device, dtype, pos_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)], pos_weight, dtype=dtype, device=device),
+            ).shape
+            == ()
+        )
+
     def test_jit(self, device, dtype):
         logits = torch.rand(2, 3, 2, dtype=dtype, device=device)
         labels = torch.rand(2, 3, 2, dtype=dtype, device=device)