Hey,
Does anyone know how I can implement a Tversky loss in a multiclass segmentation problem (I have 4 classes)? I saw this on github Multi class · Issue #3 · nabsabraham/focal-tversky-unet · GitHub 2, I change the code to pytorch, but did not understand this permutation. And when I try in my model I got this following error: “RuntimeError: number of dims don’t match in permute”.
My code:
class TverskyLoss(nn.Module):
    def __init__(self, alpha_t=0.5, beta_t=0.5,
    #weight=True, size_average=True,
    ):
        super(TverskyLoss, self).__init__()
        self.alpha_t = alpha_t
        self.beta_t = beta_t
        

    def forward(self, label, pred, smooth=1):
        # comment out if your model contains a sigmoid or equivalent activation layer
        #label = torch.sigmoid(label)

        # flatten label and prediction tensors

        label = label.permute(3,1,2,0)
        pred = pred.permute(3,1,2,0)

        flat_label = label.flatten()
        flat_pred = pred.flatten()
        TP = torch.sum(flat_label * flat_pred, 1)
        FN = torch.sum(flat_label * (1-flat_pred), 1)
        FP = torch.sum((1-flat_label)*flat_pred, 1)

        Tversky = (TP + smooth) / (TP + self.alpha_t * FP + self.beta_t * FN + smooth)
        print(self.alpha_t, self.beta_t)

        return 1 - Tversky

Someone could help me?