# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/ml.loss.ipynb.

# %% auto 0
__all__ = ['BinaryFocalLoss']

# %% ../../nbs/ml.loss.ipynb 3
import torch
from torch import nn
from fastcore.basics import ifnone

# %% ../../nbs/ml.loss.ipynb 4
# license: https://www.kaggle.com/crissallan/pytorchlightning-efficientnet-focalloss-training
class BinaryFocalLoss(nn.Module):
    """
    <sup>
    T.-Y. Lin, P. Goyal, R. Girshick, K. He, and P. Dollar, “Focal loss for dense object detection,” in Proceedings of the IEEE International Conference on Computer Vision (ICCV), Oct 2017.
    </sup>
    """
    def __init__(self, alpha=1, gamma=2, weight=None, reduction=torch.mean, thresh=0.5):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.weight = None if weight is None else torch.tensor([weight], dtype=torch.float, requires_grad=False, device='cuda')
        self.reduction = reduction
        self.thresh = 0.5
        self.eps = 1e-12 # prevent training from Nan-loss error

    def forward(self, logits, targets):
        """
        logits & target should be tensors with shape [batch_size, num_classes]
        """
        probs = torch.sigmoid(logits)
        log_pt = targets * torch.log(probs + self.eps) + (1.0 - targets) * torch.log(1.0 - probs + self.eps)
        focal_loss = -self.alpha * ((1 - torch.exp(log_pt)) ** self.gamma) * log_pt
        if self.weight is not None:
            focal_loss = focal_loss * self.weight
        
        reduction = ifnone(self.reduction, lambda x:x)
        return reduction(focal_loss)
    
    def decodes(self, x):    
        return x > self.thresh

    def activation(self, x): 
        return torch.sigmoid(x)
