fix error in tversky-loss

Author: Hsuxu

TverskyLoss/binarytverskyloss.py

@@ -102,7 +102,7 @@ def __init__(self, alpha=0.3, beta=0.7, ignore_index=None, reduction='mean'):
         """Dice loss of binary class
         Args:
             alpha: controls the penalty for false positives.
-            beta: penalty for false negative.
+            beta: penalty for false negative. Larger beta weigh recall higher
             ignore_index: Specifies a target value that is ignored and does not contribute to the input gradient
             reduction: Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'
         Shapes:
@@ -117,29 +117,31 @@ def __init__(self, alpha=0.3, beta=0.7, ignore_index=None, reduction='mean'):
         self.alpha = alpha
         self.beta = beta
         self.ignore_index = ignore_index
-        self.epsilon = 1e-6
+        self.smooth = 10
         self.reduction = reduction
         s = self.beta + self.alpha
-        if sum != 1:
+        if s != 1:
             self.beta = self.beta / s
             self.alpha = self.alpha / s
 
-    def forward(self, output, target):
+    def forward(self, output, target, mask=None):
         batch_size = output.size(0)
-
+        bg_target = 1 - target
         if self.ignore_index is not None:
             valid_mask = (target != self.ignore_index).float()
             output = output.float().mul(valid_mask)  # can not use inplace for bp
             target = target.float().mul(valid_mask)
+            bg_target = bg_target.float().mul(valid_mask)
 
         output = torch.sigmoid(output).view(batch_size, -1)
         target = target.view(batch_size, -1)
+        bg_target = bg_target.view(batch_size, -1)
 
         P_G = torch.sum(output * target, 1)  # TP
-        P_NG = torch.sum(output * (1 - target), 1)  # FP
+        P_NG = torch.sum(output * bg_target, 1)  # FP
         NP_G = torch.sum((1 - output) * target, 1)  # FN
 
-        tversky_index = P_G / (P_G + self.alpha * P_NG + self.beta * NP_G + self.epsilon)
+        tversky_index = P_G / (P_G + self.alpha * P_NG + self.beta * NP_G + self.smooth)
 
         loss = 1. - tversky_index
         # target_area = torch.sum(target_label, 1)