first pass fix at autograd fix

Author: brechmos

TverskyLoss/binarytverskyloss.py

@@ -8,32 +8,15 @@
 
 class FocalBinaryTverskyLoss(Function):
 
-    def __init__(ctx, alpha=0.5, beta=0.5, gamma=1.0, reduction='mean'):
-        """
-
-        :param alpha: controls the penalty for false positives.
-        :param beta: penalty for false negative.
-        :param gamma : focal coefficient range[1,3]
-        :param reduction: return mode
-        Notes:
-        alpha = beta = 0.5 -> dice coeff
-        alpha = beta = 1 -> tanimoto coeff
-        alpha + beta = 1 -> F beta coeff
-        add focal index -> loss=(1-T_index)**(1/gamma)
-
-        """
-        ctx.alpha = alpha
-        ctx.beta = beta
-        ctx.epsilon = 1e-6
-        ctx.reduction = reduction
-        ctx.gamma = gamma
-        s = ctx.beta + ctx.alpha
-        if sum != 1:
-            ctx.beta = ctx.beta / s
-            ctx.alpha = ctx.alpha / s
 
-    # @staticmethod
+    @staticmethod
     def forward(ctx, input, target):
+        _alpha = 0.5
+        _beta = 0.5
+        _gamma = 1.0
+        _epsilon = 1e-6
+        _reduction = 'mean'
+
         batch_size = input.size(0)
         _, input_label = input.max(1)
 
@@ -49,13 +32,13 @@ def forward(ctx, input, target):
         ctx.P_NG = torch.sum(input_label * (1 - target_label), 1)  # FP
         ctx.NP_G = torch.sum((1 - input_label) * target_label, 1)  # FN
 
-        index = ctx.P_G / (ctx.P_G + ctx.alpha * ctx.P_NG + ctx.beta * ctx.NP_G + ctx.epsilon)
-        loss = torch.pow((1 - index), 1 / ctx.gamma)
+        index = ctx.P_G / (ctx.P_G + _alpha * ctx.P_NG + _beta * ctx.NP_G + _epsilon)
+        loss = torch.pow((1 - index), 1 / _gamma)
         # target_area = torch.sum(target_label, 1)
         # loss[target_area == 0] = 0
-        if ctx.reduction == 'none':
+        if _reduction == 'none':
             loss = loss
-        elif ctx.reduction == 'sum':
+        elif _reduction == 'sum':
             loss = torch.sum(loss)
         else:
             loss = torch.mean(loss)
@@ -72,24 +55,30 @@ def backward(ctx, grad_out):
                         = 2*P_G
         (dT_loss/d_p0)=
         """
+        _alpha = 0.5
+        _beta = 0.5
+        _gamma = 1.0
+        _reduction = 'mean'
+        _epsilon = 1e-6
+
         inputs, target = ctx.saved_tensors
         inputs = inputs.float()
         target = target.float()
         batch_size = inputs.size(0)
-        sum = ctx.P_G + ctx.alpha * ctx.P_NG + ctx.beta * ctx.NP_G + ctx.epsilon
+        sum = ctx.P_G + _alpha * ctx.P_NG + _beta * ctx.NP_G + _epsilon
         P_G = ctx.P_G.view(batch_size, 1, 1, 1, 1)
         if inputs.dim() == 5:
             sum = sum.view(batch_size, 1, 1, 1, 1)
         elif inputs.dim() == 4:
             sum = sum.view(batch_size, 1, 1, 1)
             P_G = ctx.P_G.view(batch_size, 1, 1, 1)
-        sub = (ctx.alpha * (1 - target) + target) * P_G
+        sub = (_alpha * (1 - target) + target) * P_G
 
-        dL_dT = (1 / ctx.gamma) * torch.pow((P_G / sum), (1 / ctx.gamma - 1))
+        dL_dT = (1 / _gamma) * torch.pow((P_G / sum), (1 / _gamma - 1))
         dT_dp0 = -2 * (target / sum - sub / sum / sum)
         dL_dp0 = dL_dT * dT_dp0
 
-        dT_dp1 = ctx.beta * (1 - target) * P_G / sum / sum
+        dT_dp1 = _beta * (1 - target) * P_G / sum / sum
         dL_dp1 = dL_dT * dT_dp1
         grad_input = torch.cat((dL_dp1, dL_dp0), dim=1)
         # grad_input = torch.cat((grad_out.item() * dL_dp0, dL_dp0 * grad_out.item()), dim=1)
@@ -152,4 +141,4 @@ def forward(self, output, target, mask=None):
             loss = torch.sum(loss)
         else:
             loss = torch.mean(loss)
-        return loss
+        return loss

TverskyLoss/multitverskyloss.py

@@ -38,7 +38,7 @@ def forward(self, inputs, targets):
             input_idx = torch.cat((1 - input_idx, input_idx), dim=1)
             target_idx = (targets == idx) * 1
             loss_func = FocalBinaryTverskyLoss(self.alpha, self.beta, self.gamma)
-            loss_idx = loss_func(input_idx, target_idx)
+            loss_idx = loss_func.apply(input_idx, target_idx)
             weight_losses+=loss_idx * weights[idx]
         # loss = torch.Tensor(weight_losses)
         # loss = loss.to(inputs.device)