I am new to pytorch. I’m working on semantic segmentation, so I like to use the dice_loss to update the model’s parameters (Previously, I tested the model with the CrossEntropy loss function and works reasonably).
Here, some code:
for epoch in range(epochs):
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1) # since I will work with 1s and 0s
loss = dice_loss(predicted,labels)
loss.backward() # Backward propagation
optimizer.step()
A loss output is as like tensor(0.7531, grad_fn=<AddBackward0>), and loss.backward() does not fails, but in my test the loss never decreases after it inits.
Here my loss function in details:
def dice_loss(predicted, labels):
"""Dice coeff loss for a batch"""
# both the predicted and the labels data are being one-hot encoded
onehot_pred = torch.Tensor()
onehot_lab = torch.Tensor()
for batch, data in enumerate(zip(predicted, labels)):
# to_categorical is the KERAS adapted function
pred = utils.to_categorical(data[0]).unsqueeze(dim=0)
lab = utils.to_categorical(data[1]).unsqueeze(dim=0)
onehot_pred = torch.cat((onehot_pred,pred),dim=0)
onehot_lab = torch.cat((onehot_lab,lab),dim=0)
# calculate the loss function
ratio = 1 / predicted.size(0) # instead to divide by the batch_size
# loss accumulator
dc = torch.tensor(0).float()
# I put required gradient in order to create a grad_fn
# without loss.backward() does not work
onehot_pred.requires_grad = True
onehot_lab.requires_grad = True
for batch , data in enumerate(zip(onehot_pred,onehot_lab)):
dc += dice_coeff(data[0], data[1])*ratio
return dc
and the other function
def dice_coeff(predicted, labels):
“”“Dice coeff for a simple plane”“”
eps = 0.0001
inter = torch.dot(predicted.view(-1), labels.view(-1))
union = torch.sum(predicted) + torch.sum(labels) + eps
dice_coeff = 1 - (2 * inter.float() + eps) / union.float()
return dice_coeff