Unable to train resnet_18_adf

[SaumilShah66]

Hey,
I am trying to train resnet18_adf network from the beginning, but it doesn't train. After several epochs, my loss and accuracy don't change and the output shows "nan" values. I guess this is happening because my variances overshoot and they start showing "inf".

[mattiasegu]

Hi @SaumilShah66

Thanks for the interest in our work!

Concerning the NaNs for loss when training ADF, it is a known problem that training with the heteroscedastic loss may be difficult because of numerical instability problems. We also mentioned in our paper that it wasn't possible to train the heteroscedastic neural network from Kendall et al. because of numerical instability enhanced by the SoftMax layer. To address this problem when trainining the ADF network with the heteroscedastic loss, we initialized the network weights from the best pretrained ckpt on Resnet-18 with and without dropout. You can try it yourself, no modification to the code are needed, you only need to load one of the two available ckpts trained without ADF.