focal_loss.py

def binary_focal_loss_with_logits(input, target, gamma, reduction="mean"):
    t = target.float()
    p = torch.clamp((2 * t - 1) * torch.sigmoid(input) + (1 - t), min=1e-20)
    loss = - ( 1 - p) ** gamma * p.log()
    if reduction == "none":
        return loss
    elif reduction == "mean":
        return loss.mean()
    elif reduction == "sum":
        return loss.sum()
    else:
        raise ValueError("Unkown reduction {}".format(reduction))