Getting cls_logits NaN or Inf during training

[AceMcAwesome77]

I am training this retinanet 3D detection model with mostly the same parameters as the example in this repo, except with batch_size in the config = 1 because many image volumes are smaller than the training patch size. During training, I am getting this error at random, several epochs into the training:

Traceback of TorchScript, original code (most recent call last):
File "/home/mycomputer/.local/lib/python3.10/site-packages/monai/apps/detection/networks/retinanet_network.py", line 130, in forward
if torch.isnan(cls_logits).any() or torch.isinf(cls_logits).any():
if torch.is_grad_enabled():
raise ValueError("cls_logits is NaN or Inf.")
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
else:
warnings.warn("cls_logits is NaN or Inf.")
builtins.ValueError: cls_logits is NaN or Inf.

On the last few training attempts, this failed on epoch 6 on the first two attempts, then failed on epoch 12 on the third attempt. So it can make it though all the training data without failing on any particular case. Does anyone know what could be causing this? If it's exploding gradients, is there something built into MONAI to clip these and prevent the training from crashing? Thanks!

[KumoLiu]

Hi @AceMcAwesome77,

The error message you're encountering, "cls_logits is NaN or Inf.", is letting you know that at some point in your training, the cls_logits tensor contains a Not a Number (NaN) or Infinity (Inf).
This can occur due to various reasons. This could come from a learning rate that's too high, instabilities in your numerical operations, uninitialized variables, or it could also be a problem with the specific data you're inputting into the model. It's a sign that the model is diverging, and gradients are getting out of control, which can also stem from exploding or vanishing gradients.
You indeed can attempt to mitigate this issue using gradient clipping, which can help ensure gradients never exceed a certain threshold. However, applying gradient clipping doesn't guarantee to resolve the root cause of the problem.
I would recommend looking at your training process more holistically. Inspect the learning rate, look for possible issues in the data, try normalizing the inputs, or use different weight initialization techniques.

Hope it helps, thanks.

[AceMcAwesome77]

Thanks for the reply and advice! I tried lowering the learning rate quite a bit but it still errors in the same way. Since the NaN loss does not occur until long into training, in some cases past epoch 40, it doesn't seem like it's caused by the data itself. I added logging that prints out the specific training sample it errors on, and it's a different sample each time, but interestingly, they are always negative samples, i.e. samples with no ground truth bounding boxes. My dataset has about a 3:1 ratio of negatives to positives but I looked at 20 errored samples and they were all negatives. This leads me to think the error has to do specifically with not handling cases with no ground truth boxes well. I know this functionality was added to the newer MONAI versions since back in version 1.1.0 it had to be "hacked" in as seen in this issue:
#1292
Do you have an idea for why the recent inclusion of the StandardizeEmptyBoxd function and the general empty box logic throughout the code might cause NaN loss as training goes along? If I wanted to edit the class loss to something hacky like "cls_loss = 0 when no gt boxes if predicted no boxes", and "cls_loss = 1 when no gt boxes if predicted a box", where in the code would I do that? Thanks!

[AceMcAwesome77]

I have managed to improve this issue by using a train batch size of 8 (it had been 1 before) and using an ADAM optimizer with learning rate = 10^-5. It still goes to NaN loss occasionally but can make it over 100 epochs before doing so. If it errors I re-initialize the optimizer and keep going and it continues training fairly well. I think the small batch size was contributing to instability.

[ericspod]

Larger batch sizes do often help especially with BatchNorm or other layers sensitive to number of inputs. The NaN values on occasion though are still concerning but hard to pin down as @KumoLiu discusses. Once common cause is input data which is all 0 sometimes due to preprocessing errors, this would cause an issue very quickly if your batch size is 1 but if it's 8 this happens when you're unlucky enough to generate 8 empty inputs. It's worth checking this isn't happening.