Is this a correct implementation for focal loss in pytorch?

mr_ruirui: <p>Hello, I am new to pytorch and currently focusing on  text classification task using deep learning networks. The dataset contains two classes and the dataset highly imbalanced(pos:neg==100:1). So I want to use focal loss to have a try.</p>
<p>I have seen some focal loss implementations but they are a little bit hard to write. So I implement the focal loss(<a href="https://arxiv.org/abs/1708.02002" rel="nofollow noopener">Focal Loss for Dense Object Detection</a>) with pytorch==1.0 and python==3.6.5. It works just the same as standard binary cross entropy loss, sometimes worse. Did I correctly implement it?</p>
<p>Here is the code:</p>
<pre><code class="lang-auto">class FocalLoss(nn.Module):
"""
binary focal loss
"""

def __init__(self, alpha=0.25, gamma=2):
    super(FocalLoss, self).__init__()
    self.weight = torch.Tensor([alpha, 1-alpha])
    self.nllLoss = nn.NLLLoss(weight=self.weight)
    self.gamma = gamma

def forward(self, input, target):
    softmax = F.softmax(input, dim=1)
    log_logits = torch.log(softmax)
    fix_weights = (1 - softmax) ** self.gamma
    logits = fix_weights * log_logits
    return self.nllLoss(logits, target)
</code></pre>

n0obcoder: <p>i am also in interested in knowing the implementation of focal loss.it would be great if someone who has implemented it, could help us out here !</p>

Diego: <p>You can checkout kornia’s implementation <a href="https://github.com/arraiyopensource/kornia/blob/master/kornia/losses/focal.py" rel="nofollow noopener">here</a></p>

n0obcoder: <p>its input has shape (batch_size, num_classes, H, W). but i have an input matrix of shape (batch_size, num_classes) and target matrix of same shape (batch_size, num_classes). So i have come up with this…</p>
<pre><code class="lang-auto">
import torch, sys, os, pdb
import torch.nn as nn
import torch.nn.functional as F
</code></pre>
<pre><code class="lang-auto">class FocalLoss(nn.Module):

    def __init__(self, gamma = 1.0):
        super(FocalLoss, self).__init__()
        self.gamma = torch.tensor(gamma, dtype = torch.float32)
        self.eps = 1e-6

    def forward(self, input, target):
        # input are not the probabilities, they are just the cnn out vector
        # input and target shape: (bs, n_classes)
        # sigmoid
        probs = torch.sigmoid(input)
        log_probs = -torch.log(probs)

        focal_loss = torch.sum(  torch.pow(1-probs + self.eps, self.gamma).mul(log_probs).mul(target)  , dim=1)
        # bce_loss = torch.sum(log_probs.mul(target), dim = 1)
        
        return focal_loss.mean() #, bce_loss
</code></pre>
<p>pls let me know if its good or not</p>

Diego: <p>Try this:</p>
<pre><code class="lang-auto">BCE_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
pt = torch.exp(-BCE_loss) # prevents nans when probability 0
F_loss = self.alpha * (1-pt)**self.gamma * BCE_loss
return focal_loss.mean()
</code></pre>
<p>Remember the alpha to address class imbalance and keep in mind that this will only work for binary classification.<br>
I think this is very similar to your implementation it just uses the BCE function which does the sigmoid and those <code>.mul()</code> for you and also prevents the <code>nan</code> problem that is present in your implementation when probability is 0 i.e log(0) = -inf<br>
There are also some very nice implementations that work for multiclass problems. These implementations can be found <a href="https://github.com/kuangliu/pytorch-retinanet/blob/2d7c663350f330a34771a8fa6a4f37a2baa52a1d/loss.py#L29" rel="nofollow noopener">here</a> and <a href="https://github.com/DingKe/pytorch_workplace/blob/4bc60a2c3640de522d0b72262667ba70391ba16e/focalloss/loss.py#L22" rel="nofollow noopener">here</a></p>

achaiah: <p>Hey Diego,</p>
<p>I have a couple of questions about your implementation.</p>
<aside class="quote no-group" data-username="Diego" data-post="5" data-topic="43327">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/diego/48/6989_2.png" class="avatar"> Diego:</div>
<blockquote>
<p>F_loss = self.alpha * (1-pt)**self.gamma * BCE_loss</p>
</blockquote>
</aside>
<ol>
<li>What is alpha for a binary focal loss?  In binary segmentation there is only one class (the one you’re trying to segment)</li>
<li>Your equation isn’t using epsilon, is that intentional?</li>
</ol>
<p>Thanks!</p>

Diego: <p>Hi, just to clarify this is not my implementation, I found it somewhere in kaggle, I cant find the link now. I will try to answer your questions nonetheless:</p>
<ol>
<li>Alpha is hyperparameter that you can tune to assign more importance to samples from class A or B. I dont know anything about binary segmentation so correct me if i am wrong, but I assume that it must have at least 2 classes. The class you are trying to segment and background.</li>
<li>The epsilon is used to avoid numerical instability if probability equals 0. In this case  <code>torch.exp()</code> will deal with that.</li>
</ol>
<p>You can find another, perhaps more clear, implementation <a href="https://discuss.pytorch.org/t/focal-loss-for-binary-semantic-segmentation-implementation/12125">here</a></p>

arkrde: <p>Hi,<br>
Here is my implementation. I have tried to use the info on <code>torch.nn.functional.log_softmax</code> from <a href="https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.log_softmax" rel="nofollow noopener">https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.log_softmax</a>. So far, it’s working well in a class - imbalance problem.<br>
Please let me know if it works for you.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

class FocalLoss(nn.Module):
    
    def __init__(self, weight=None, 
                 gamma=2., reduction='none'):
        nn.Module.__init__(self)
        self.weight = weight
        self.gamma = gamma
        self.reduction = reduction
        
    def forward(self, input_tensor, target_tensor):
        log_prob = F.log_softmax(input_tensor, dim=-1)
        prob = torch.exp(log_prob)
        return F.nll_loss(
            ((1 - prob) ** self.gamma) * log_prob, 
            target_tensor, 
            weight=self.weight,
            reduction = self.reduction
        )
</code></pre>

pigfarm: <p>Looks working for me too. My Input size is NxC, and target size is N.</p>

coyote: <p>Hey, thanks for the implementation. Could you tell something about weights variable ? What do you use for that ?</p>

arkrde: <p>Weights should be a  1-d tensor indicating the relative class importance. For a balanced case i.e. <code>weight=None</code>, it’s equivalent to a 1-d tensor whose values are all equal e.g. <code>1</code>. For class-imbalance problems, this can be tweaked to adjust for the imbalance i.e. <code>[0.5, 1]</code> in a binary classification problem where the first class is twice more likely to appear than the second in the target variable. Take a look at the mathematical formulations here: <a href="https://pytorch.org/docs/stable/nn.html#nllloss" rel="nofollow noopener">https://pytorch.org/docs/stable/nn.html#nllloss</a></p>

calusbr: <p>Don’t you use alpha in your implementation?</p>

BaruchYoussin: <p>kornia’s implementation does not seem to treat alpha properly.<br>
It seems to use the same value for both classes.</p>

BaruchYoussin: <p>Here is my implementation along the lines of <a class="mention" href="/u/diego">@Diego</a> with alpha implemented.  Binary case only.</p>
<pre><code class="lang-auto">def focal_loss(bce_loss, targets, gamma, alpha):
    """Binary focal loss, mean.

    Per https://discuss.pytorch.org/t/is-this-a-correct-implementation-for-focal-loss-in-pytorch/43327/5 with
    improvements for alpha.
    :param bce_loss: Binary Cross Entropy loss, a torch tensor.
    :param targets: a torch tensor containing the ground truth, 0s and 1s.
    :param gamma: focal loss power parameter, a float scalar.
    :param alpha: weight of the class indicated by 1, a float scalar.
    """
    p_t = torch.exp(-bce_loss)
    alpha_tensor = (1 - alpha) + targets * (2 * alpha - 1)  # alpha if target = 1 and 1 - alpha if target = 0
    f_loss = alpha_tensor * (1 - p_t) ** gamma * bce_loss
    return f_loss.mean()
</code></pre>

gokul_prasad: <p>I believe alpha and the weight parameter is essentially doing the same thing here. they are just linearly scaling the loss. This is a neat way of implementing alpha. <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>

gokul_prasad: <p>Here is an implementation of Focal Loss for muti-class classification:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/f/7/f7df80a29ae7971c6ee551ed6b26fbdc44fa50e9.png" alt="1_gO_nxGFmpAelOrU_D9O5-Q" data-base62-sha1="zmMMNs2Gpj7rWbABlM45tC9zW6l" width="234" height="26"></p>
<p>Here, -log(pt) is our ordinary cross entropy loss.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

class FocalLoss(nn.modules.loss._WeightedLoss):
    def __init__(self, weight=None, gamma=2,reduction='mean'):
        super(FocalLoss, self).__init__(weight,reduction=reduction)
        self.gamma = gamma
        self.weight = weight #weight parameter will act as the alpha parameter to balance class weights

    def forward(self, input, target):

        ce_loss = F.cross_entropy(input, target,reduction=self.reduction,weight=self.weight) 
        pt = torch.exp(-ce_loss)
        focal_loss = ((1 - pt) ** self.gamma * ce_loss).mean()
        return focal_loss
</code></pre>

FaultyBagnose: <p>Does this work properly for minibatches?<br>
I guess F.cross_entropy() gives the average c-e entropy over the batch, and pt is a scalar variable that modifies the loss for the batch. So, if some of the input-target patterns have a low and some have a high ce_loss they get the same focal adjustment?<br>
If so, this might fix it:</p>
<pre><code>def forward(self, input, target):
    patterns = target.shape[0]
    tot = 0
    for b in range(patterns):
        ce_loss = F.cross_entropy(input[b:b+1,], target[b:b+1],reduction=self.reduction,weight=self.weight)
        pt = torch.exp(-ce_loss)
        focal_loss = (1 - pt) ** self.gamma * ce_loss
        tot = tot + focal_loss
    return tot/patterns</code></pre>

Samue1: <p>Not sure, but is <code>pt</code> not the network’s prediction for class <code>t</code>?</p>

theGuyWithBlackTie: <p><code>F.cross_entropy</code> takes logits from the model. Logits are outputs of the model, they are not probabilities. That’s the reason, for probabilities (i.e. pt), <code>torch.exp(-ce_loss)</code> is done.</p>
<p>Hope this helps.</p>

Songhua_Hu: <p>This may not be correct if you have class weight, since the weight is attached before multiplying by (1-p)**gamma. I think <a class="mention" href="/u/arkrde">@arkrde</a> 's answer works. In pytorch, F.cross_entropy is equal to F.nll_loss(F.log_softmax). The weight (i.e. the alpha) should be attached on nll_loss().</p>