Fix FocalLoss (#259)

* Fix positional arguments for FocalLoss * Set loss_0 (1 - alpha) coeff in binary_focal_loss
qubvel · Nov 18, 2019 · d71189c · d71189c
1 parent ba5cfca
commit d71189c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/segmentation_models/base/functional.py b/segmentation_models/base/functional.py
@@ -303,6 +303,6 @@ def binary_focal_loss(gt, pr, gamma=2.0, alpha=0.25, **kwargs):
     pr = backend.clip(pr, backend.epsilon(), 1.0 - backend.epsilon())
 
     loss_1 = - gt * (alpha * backend.pow((1 - pr), gamma) * backend.log(pr))
-    loss_0 = - (1 - gt) * (alpha * backend.pow((pr), gamma) * backend.log(1 - pr))
+    loss_0 = - (1 - gt) * ((1 - alpha) * backend.pow((pr), gamma) * backend.log(1 - pr))
     loss = backend.mean(loss_0 + loss_1)
     return loss
diff --git a/segmentation_models/losses.py b/segmentation_models/losses.py
@@ -200,8 +200,8 @@ def __call__(self, gt, pr):
         return F.categorical_focal_loss(
             gt,
             pr,
-            self.alpha,
-            self.gamma,
+            alpha=self.alpha,
+            gamma=self.gamma,
             class_indexes=self.class_indexes,
             **self.submodules
         )
@@ -235,7 +235,7 @@ def __init__(self, alpha=0.25, gamma=2.):
         self.gamma = gamma
 
     def __call__(self, gt, pr):
-        return F.binary_focal_loss(gt, pr, self.alpha, self.gamma, **self.submodules)
+        return F.binary_focal_loss(gt, pr, alpha=self.alpha, gamma=self.gamma, **self.submodules)
 
 
 # aliases