File size: 850 Bytes
9e14838 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | import torch
def focal_loss(input, target, gamma=2, weight=None):
# Code based on: https://github.com/clcarwin/focal_loss_pytorch/blob/master/focalloss.py
assert len(input.shape) == 2
assert len(target.shape) == 1
# input.shape = N,C
# target.shape = N
target = target.unsqueeze(1)
logpt = torch.log_softmax(input, dim=-1)
logpt = torch.clamp(logpt, max=-0.01, min=-5)
logpt = logpt.gather(1, target)
logpt = logpt.view(-1)
pt = torch.tensor(logpt.data.exp())
if (not weight is None):
# target.shape = N, 1
weight = torch.tensor(
weight,
device=input.device,
dtype=input.dtype
)
at = weight.gather(0, target.data.view(-1))
else:
at = 1
logpt = logpt * at
loss = -1 * (1 - pt)**gamma * logpt
return loss
|