Hello all, I am using dice loss for multiple class (4 classes problem). I want to use weight for each class at each pixel level. So, my weight will have size of BxCxHxW (C=4) in my case. How can I use the weight to assign to dice loss? This is my current solution that multiple the weight with the input (network prediction) after softmax
class SoftDiceLoss(nn.Module):
def __init__(self, n_classes):
super(SoftDiceLoss, self).__init__()
self.one_hot_encoder = One_Hot(n_classes).forward
self.n_classes = n_classes
def forward(self, input, target, weight):
smooth = 0.01
batch_size = input.size(0)
input = F.softmax(input, dim=1)
input = input*weight
input = input.view(batch_size, self.n_classes, -1)
target = self.one_hot_encoder(target).contiguous().view(batch_size, self.n_classes, -1)
inter = torch.sum(input * target, 2) + smooth
union = torch.sum(input, 2) + torch.sum(target, 2) + smooth
score = torch.sum(2.0 * inter / union)
score = 1.0 - score / (float(batch_size) * float(self.n_classes))
return score
And the second solution is that multiply the weight in the inter and union position
class SoftDiceLoss(nn.Module):
def __init__(self, n_classes):
super(SoftDiceLoss, self).__init__()
self.one_hot_encoder = One_Hot(n_classes).forward
self.n_classes = n_classes
def forward(self, input, target, weight):
smooth = 0.01
batch_size = input.size(0)
input = F.softmax(input, dim=1).view(batch_size, self.n_classes, -1)
target = self.one_hot_encoder(target).contiguous().view(batch_size, self.n_classes, -1)
weight = weight.view(batch_size, self.n_classes, -1)
inter = torch.sum(input * target * weight, 2) + smooth
union = torch.sum(input*weight, 2) + torch.sum(target*weight, 2) + smooth
score = torch.sum(2.0 * inter / union)
score = 1.0 - score / (float(batch_size) * float(self.n_classes))
return score
Which one is correct?