Handling exploding gradients

[rubeea]

Hi,

I am training my model with PSP decode head and on my custom dataset which is highly imbalanced. I am using dice loss due to the imbalanced nature of the dataset but my gradients are exploding and are in the order of ~5000s. I am using Adam optimizer with LR 3e-4 and weight decay 0.0001 as per the mmsegmentation documentation. I even used gradient clipping according to optimizer_config = dict(
delete=True, grad_clip=dict(max_norm=35, norm_type=2))
but the least value where the gradients go to is around ~3000s. Can anyone tell me how can I stabilize the training process and get my gradients down to 1.0. My targets are in the range of [0,1] (binary semantic segmentation task).
Any help will be highly appreciated.

[xvjiarui]

Hi @rubeea
How about other losses? I am not sure whether it dues to Dice loss.

[rubeea]

Hi @rubeea
How about other losses? I am not sure whether it dues to Dice loss.

Cross entropy is working fine but results in vashing gradients with longer training times. However, if I want to use a region based loss such as the Dice loss should I define the backward function in my loss implementation in addition to the forward function as well? The backward function will handle the backpropagation mechanism. Moreover, can you kindly provide an explanation of the gradient clipping parameters optimizer_config = dict(
delete=True, grad_clip=dict(max_norm=35, norm_type=2)) as the documentation page isn't available.
Thanks

[rubeea]

Hi @xvjiarui,

I figured out the problem on why my gradients were too high. I was applying weighted loss reduction on the dice loss with the multiplication of weights.sum() and the dice loss and hence the gradients were too high. My gradients are now stable with just reduce_loss function. However, the model does not give good results on the minority class with these region based losses (dice loss), although as per the implementation papers, dice loss is supposed to work well for imbalanced datasets. The gradients from the dice loss are of the order of ~0.05 no matter how longer I train. Any suggestion on how I can improve these region based loss functions for my imbalanced dataset implementation?

Any help in this regard will be highly appreciated
Thanks