Softmax focal loss references: Focal Loss for Dense Object Detection https://github.com/Hsuxu/FocalLoss-PyTorch
| 47 | |
| 48 | |
| 49 | class FocalLoss(nn.Module): |
| 50 | """Softmax focal loss |
| 51 | references: Focal Loss for Dense Object Detection |
| 52 | https://github.com/Hsuxu/FocalLoss-PyTorch |
| 53 | """ |
| 54 | |
| 55 | def __init__(self, label_size, activation_type=ActivationType.SOFTMAX, |
| 56 | gamma=2.0, alpha=0.25, epsilon=1.e-9): |
| 57 | super(FocalLoss, self).__init__() |
| 58 | self.num_cls = label_size |
| 59 | self.activation_type = activation_type |
| 60 | self.gamma = gamma |
| 61 | self.alpha = alpha |
| 62 | self.epsilon = epsilon |
| 63 | |
| 64 | def forward(self, logits, target): |
| 65 | """ |
| 66 | Args: |
| 67 | logits: model's output, shape of [batch_size, num_cls] |
| 68 | target: ground truth labels, shape of [batch_size] |
| 69 | Returns: |
| 70 | shape of [batch_size] |
| 71 | """ |
| 72 | if self.activation_type == ActivationType.SOFTMAX: |
| 73 | idx = target.view(-1, 1).long() |
| 74 | one_hot_key = torch.zeros(idx.size(0), self.num_cls, |
| 75 | dtype=torch.float, |
| 76 | device=idx.device) |
| 77 | one_hot_key = one_hot_key.scatter_(1, idx, 1) |
| 78 | logits = torch.softmax(logits, dim=-1) |
| 79 | loss = -self.alpha * one_hot_key * \ |
| 80 | torch.pow((1 - logits), self.gamma) * \ |
| 81 | (logits + self.epsilon).log() |
| 82 | loss = loss.sum(1) |
| 83 | elif self.activation_type == ActivationType.SIGMOID: |
| 84 | multi_hot_key = target |
| 85 | logits = torch.sigmoid(logits) |
| 86 | zero_hot_key = 1 - multi_hot_key |
| 87 | loss = -self.alpha * multi_hot_key * \ |
| 88 | torch.pow((1 - logits), self.gamma) * \ |
| 89 | (logits + self.epsilon).log() |
| 90 | loss += -(1 - self.alpha) * zero_hot_key * \ |
| 91 | torch.pow(logits, self.gamma) * \ |
| 92 | (1 - logits + self.epsilon).log() |
| 93 | else: |
| 94 | raise TypeError("Unknown activation type: " + self.activation_type |
| 95 | + "Supported activation types: " + |
| 96 | ActivationType.str()) |
| 97 | return loss.mean() |
| 98 | |
| 99 | |
| 100 | class ClassificationLoss(torch.nn.Module): |