| 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 | |