Hello All,
I am trying to implement develop a model to generate masks on the Kvasir-SEG Dataset(The Kvasir-SEG Dataset) and using a mixture of focal and Dice loss which is as follows -
class DiceFocalLoss(nn.Module):
def __init__(self, weight=None, size_average=True, alpha=0.25, gamma=3, smooth=1):
super(DiceFocalLoss, self).__init__()
self.gamma = gamma
self.alpha = alpha
self.smooth = smooth
def forward(self, inputs, targets):
# comment out if your model contains a sigmoid or equivalent activation layer
#inputs = F.sigmoid(inputs)
# flatten label and prediction tensors
inputs = inputs.view(-1)
targets = targets.view(-1)
intersection = (inputs * targets).sum()
dice_loss = 1 - (2. * intersection + self.smooth) / (inputs.sum() + targets.sum() + self.smooth)
BCE = F.binary_cross_entropy_with_logits(inputs, targets, reduction='mean')
BCE_EXP = torch.exp(-BCE)
focal_loss = self.alpha * (1-BCE_EXP)**self.gamma * BCE
Dice_BCE = focal_loss + dice_loss
return Dice_BCE
and while debugging the model, I take a single (image, mask) pair from my dataset and train only on it but surprisingly while using autocast the loss change but rather remains constant indicating the model weights are not being updated. I am able to recreate the issue with a simple 2 layer convolution model as well -
the toy model is defined as -
model = nn.Sequential(nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, stride=1, padding=1),
nn.Conv2d(in_channels=3, out_channels=1, kernel_size=1, stride=1, padding=0))
model = model.cuda().train()
and my training loop is -
optim = torch.optim.SGD(model.parameters(), lr=1e-5)
for i in range(10):
with autocast():
out = model(image)
l = loss(out, mask)
scaler.scale(l).backward()
scaler.step(optim)
scaler.update()
print(l)
loss values without autocast -
tensor(2.6734, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.0208, device='cuda:0', grad_fn=<AddBackward0>)
tensor(1.3117, device='cuda:0', grad_fn=<AddBackward0>)
tensor(0.8833, device='cuda:0', grad_fn=<AddBackward0>)
tensor(0.7073, device='cuda:0', grad_fn=<AddBackward0>)
tensor(0.6678, device='cuda:0', grad_fn=<AddBackward0>)
tensor(0.6844, device='cuda:0', grad_fn=<AddBackward0>)
tensor(0.7187, device='cuda:0', grad_fn=<AddBackward0>)
tensor(0.7563, device='cuda:0', grad_fn=<AddBackward0>)
tensor(0.7941, device='cuda:0', grad_fn=<AddBackward0>)
Loss values with autocast -
tensor(8.9766, device='cuda:0', grad_fn=<AddBackward0>)
tensor(8.9766, device='cuda:0', grad_fn=<AddBackward0>)
tensor(8.9766, device='cuda:0', grad_fn=<AddBackward0>)
tensor(8.9766, device='cuda:0', grad_fn=<AddBackward0>)
tensor(8.9766, device='cuda:0', grad_fn=<AddBackward0>)
tensor(8.9766, device='cuda:0', grad_fn=<AddBackward0>)
tensor(8.9766, device='cuda:0', grad_fn=<AddBackward0>)
tensor(8.9766, device='cuda:0', grad_fn=<AddBackward0>)
tensor(8.9766, device='cuda:0', grad_fn=<AddBackward0>)
tensor(8.9766, device='cuda:0', grad_fn=<AddBackward0>)
As you can see, there is no change no matter what which implies the model weights are not getting updated.
Could anyone please point out what is going wrong ?
TIA
. On my single data point, the model converges decently in both modes.