(self, gamma=2, weight=None, size_average=True, ignore_index=255)
| 49 | |
| 50 | class FocalLoss2d(nn.Module): |
| 51 | def __init__(self, gamma=2, weight=None, size_average=True, ignore_index=255): |
| 52 | super(FocalLoss2d, self).__init__() |
| 53 | self.gamma = gamma |
| 54 | self.nll_loss = nn.NLLLoss2d(weight, size_average, ignore_index) |
| 55 | |
| 56 | def forward(self, inputs, targets): |
| 57 | return self.nll_loss((1 - F.softmax(inputs)) ** self.gamma * F.log_softmax(inputs), targets) |