Hi All,
I am trying to implement dice loss for semantic segmentation using FCN_resnet101. For some reason, the dice loss is not changing and the model is not updated.
import torch
import torchvision
import loader
from loader import DataLoaderSegmentation
import torch.nn as nn
import torch.optim as optim
import numpy as np
from torch.utils.data.sampler import SubsetRandomSampler
from torch.autograd import Variable

batch_size = 1
validation_split = .2
shuffle_dataset = True
random_seed= 66

n_class    = 2
num_epochs = 5
lr         = 1e-4
momentum   = 0.9
w_decay    = 1e-5
step_size  = 50
gamma      = 0.5
cw = [0.5028138887905369,89.34501791284346]
class_weight=torch.FloatTensor(cw).cuda()
traindata = DataLoaderSegmentation('/home/ubuntu/Downloads/imgs/lensonly')
trainloader = torch.utils.data.DataLoader(traindata)

class diceloss(torch.nn.Module):
    def init(self):
        super(diceLoss, self).init()
    def forward(self,pred, target):
       smooth = 1.
       iflat = pred.contiguous().view(-1)
       tflat = target.contiguous().view(-1)
       intersection = (iflat * tflat).sum()
       A_sum = torch.sum(iflat * iflat)
       B_sum = torch.sum(tflat * tflat)
       return 1 - ((2. * intersection + smooth) / (A_sum + B_sum + smooth) )

model =torchvision.models.segmentation.fcn_resnet101(pretrained=False,aux_loss=None,num_classes=2).cuda()
criterion = diceloss()

for iter in range(num_epochs):
    for (i,l) in trainloader:
        optimizer.zero_grad()
        i= i.to(device)
        l = l.to(device)
        l[l!=0]=1
        l=l.long()
        outt = model(i)
        pred = Variable(torch.argmax(outt['out'],1).float(),requires_grad=True).cuda()
        loss = criterion(pred,l.squeeze(0).type(torch.cuda.FloatTensor))
        loss.backward()
        optimizer.step()
    print(iter)
torch.save(model, '/home/ubuntu/Downloads/newnet.pth')

Any and all help is appreciated. I copy pasted dice loss from here

  
      gist.github.com
  
  
    https://gist.github.com/weiliu620/52d140b22685cf9552da4899e2160183 1.5k
Dice_coeff_loss.py
def dice_loss(pred, target):
    """This definition generalize to real valued pred and target vector.
This should be differentiable.

    pred: tensor with first dimension as batch
    target: tensor with first dimension as batch

    """

    smooth = 1.
This file has been truncated. show original




  
  
    
    
  
  


Please let me know if this is the correct way of implementing it.
Thanks
Nishanth