Args: logits: model's output, shape of [batch_size, num_cls] target: ground truth labels, shape of [batch_size] Returns: shape of [batch_size]
(self, logits, target)
| 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): |