Merge pull request #337 from urchade/fix/mean_loss_reduction

Fix normalization factor for mean loss reduction

Author: Ingvarstep

gliner/modeling/base.py

@@ -477,7 +477,8 @@ def loss(
         all_losses = all_losses * span_mask.float()
 
         if reduction == "mean":
-            loss = all_losses.mean()
+            num_valid = span_mask.float().sum()
+            loss = all_losses.sum() / num_valid if num_valid > 0 else torch.tensor(0.0, device=scores.device)
         elif reduction == "sum":
             loss = all_losses.sum()
         else:
@@ -679,7 +680,8 @@ def loss(
         all_losses = all_losses * mask
 
         if reduction == "mean":
-            loss = all_losses.mean()
+            num_valid = mask.float().sum()
+            loss = all_losses.sum() / num_valid if num_valid > 0 else torch.tensor(0.0, device=scores.device)
         elif reduction == "sum":
             loss = all_losses.sum()
         else:
@@ -976,7 +978,8 @@ def loss(
         all_losses = all_losses * mask_label.float()
 
         if reduction == "mean":
-            loss = all_losses.mean()
+            num_valid = mask_label.float().sum()
+            loss = all_losses.sum() / num_valid if num_valid > 0 else torch.tensor(0.0, device=scores.device)
         elif reduction == "sum":
             loss = all_losses.sum()
         else:
@@ -1588,7 +1591,8 @@ def loss(
         all_losses = all_losses * mask_label.float()
 
         if reduction == "mean":
-            loss = all_losses.mean()
+            num_valid = mask_label.float().sum()
+            loss = all_losses.sum() / num_valid if num_valid > 0 else torch.tensor(0.0, device=scores.device)
         elif reduction == "sum":
             loss = all_losses.sum()
         else:
@@ -2477,19 +2481,20 @@ def loss(
         """
         all_losses = self._loss(scores, labels, alpha, gamma, prob_margin, label_smoothing, negatives)
 
-        all_losses = all_losses * (word_mask.unsqueeze(-1) * prompts_embedding_mask.unsqueeze(1)).unsqueeze(-1)
+        masked_loss = all_losses * (word_mask.unsqueeze(-1) * prompts_embedding_mask.unsqueeze(1)).unsqueeze(-1)
 
         if reduction == "mean":
-            loss = all_losses.mean()
+            num_valid = (word_mask.unsqueeze(-1) * prompts_embedding_mask.unsqueeze(1)).sum()
+            loss = masked_loss.sum() / num_valid if num_valid > 0 else torch.tensor(0.0, device=scores.device)
         elif reduction == "sum":
-            loss = all_losses.sum()
+            loss = masked_loss.sum()
         else:
             warnings.warn(
                 f"Invalid Value for config 'loss_reduction': '{reduction}' \n Supported reduction modes:"
                 f" 'none', 'mean', 'sum'. It will be used 'sum' instead.",
                 stacklevel=2,
             )
-            loss = all_losses.sum()
+            loss = masked_loss.sum()
         return loss
 
     def represent_spans(