shuxinyin / nlp-loss-pytorch Goto Github PK

def forward(self, logits, targets):
        probs = torch.sigmoid(logits)
        probs = torch.gather(probs, dim=1, index=targets.unsqueeze(1))

        targets = targets.unsqueeze(dim=1)
        pos_mask = (targets == 1).float()
        neg_mask = (targets == 0).float()

        pos_weight = pos_mask * ((1 - probs) ** self.alpha) * probs
        pos_loss = 1 - (2 * pos_weight + self.smooth) / (pos_weight + 1 + self.smooth)

        neg_weight = neg_mask * ((1 - probs) ** self.alpha) * probs
        neg_loss = 1 - (2 * neg_weight + self.smooth) / (neg_weight + self.smooth)

        loss = pos_loss + neg_loss
        loss = loss.mean()
        return loss