From 37c15df080c019fcb032b9654109c4324d9fd377 Mon Sep 17 00:00:00 2001 From: Wessam Frrag Date: Fri, 27 Sep 2024 10:56:56 +0200 Subject: [PATCH] Wessam611 focal loss kwargs (#1) Feat: enable **kwargs in FocalLoss which are being passed to Torch's cross entropy functions example key: ignore_index --- cbloss/loss.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/cbloss/loss.py b/cbloss/loss.py index 78877fd..57305f6 100644 --- a/cbloss/loss.py +++ b/cbloss/loss.py @@ -46,20 +46,21 @@ class FocalLoss(nn.Module): Focal loss value. """ - def __init__(self, num_classes, alpha=1, gamma=2, reduction="mean"): + def __init__(self, num_classes, alpha=1, gamma=2, reduction="mean", **kwargs): super(FocalLoss, self).__init__() self.num_classes = num_classes self.alpha = alpha self.gamma = gamma self.reduction = reduction + self.kwargs = kwargs def forward(self, outputs, targets): if self.num_classes == 2: # binary classification problem - ce_loss = F.binary_cross_entropy_with_logits(outputs, targets.float(), reduction="none") + ce_loss = F.binary_cross_entropy_with_logits(outputs, targets.float(), reduction="none", **self.kwargs) elif self.num_classes > 2: # multi-class classification problem - ce_loss = F.cross_entropy(outputs, targets, reduction="none") + ce_loss = F.cross_entropy(outputs, targets, reduction="none", **self.kwargs) else: raise ValueError(f"Invalid input value : {self.num_classes}. It should be equal or greater than 2.")