NaN loss after some epochs

[talhaanwarch]

After some epochs loss turned out to be NaN, while doing 3D segmentation

from  monai import transforms as T
train_transforms = T.Compose(
            [
                T.LoadImaged(keys=["image", "label"]),
                T.AddChanneld(keys=["image", "label"]),
                T.Spacingd(keys=["image", "label"],pixdim=(1.5, 1.5, 2.0),mode=("bilinear", "nearest"),),
                T.Orientationd(keys=["image", "label"], axcodes="RAS"),
                T.ScaleIntensityd(keys=["image"]),
                T.Resized(keys=["image", "label"],spatial_size=(256,256,-1)),
                T.CropForegroundd(keys=["image", "label"], source_key="image"),
                T.RandCropByPosNegLabeld(keys=["image", "label"],label_key="label",spatial_size=(64, 64, 64),pos=1,neg=1,
                    num_samples=8,image_key="image", image_threshold=0,),
                T.RandFlipd(keys=["image", "label"],spatial_axis=[0],prob=0.20,),
                T.RandFlipd(keys=["image", "label"],spatial_axis=[1],prob=0.20,),
                T.RandFlipd(keys=["image", "label"],spatial_axis=[2],prob=0.20,),
                T.RandShiftIntensityd(keys=["image"],offsets=0.10,prob=0.20,),
            ]
        )
val_transforms = T.Compose(
  [
      T.LoadImaged(keys=["image", "label"]),
      T.AddChanneld(keys=["image", "label"]),
      T.Spacingd(keys=["image", "label"],pixdim=(1.5, 1.5, 2.0),mode=("bilinear", "nearest"),),
      T.Orientationd(keys=["image", "label"], axcodes="RAS"),
      T.ScaleIntensityd(keys=["image"]),
      T.Resized(keys=["image", "label"],spatial_size=(256,256,-1)),
      T.CropForegroundd(keys=["image", "label"], source_key="image"),
      T.EnsureTyped(keys=["image", "label"]),
  ]
)

from monai.networks.layers import Norm
from monai.inferers import sliding_window_inference

class OurModel(LightningModule):
  def __init__(self):
    super(OurModel,self).__init__()
    #architecute

    self.layer = nets.UNETR(
                  in_channels=1,
                  out_channels=5,
                  img_size=(64, 64, 64),
                  feature_size=16,
                  hidden_size=768,
                  mlp_dim=3072,
                  num_heads=12,
                  pos_embed="perceptron",
                  norm_name="instance",
                  res_block=True,
                  dropout_rate=0.0,
              )
  #parameters
    self.lr=1e-3
    self.bs=6
    self.numworker=1
    self.criterion=losses.DiceFocalLoss(to_onehot_y=True, softmax=True)
    self.post_pred = T.Compose([T.EnsureType(), T.AsDiscrete(argmax=True, to_onehot=True, n_classes=5)])
    self.post_label = T.Compose([T.EnsureType(), T.AsDiscrete(to_onehot=True, n_classes=5)])


  def forward(self,x):
    return self.layer(x)


  def configure_optimizers(self):
    opt=torch.optim.AdamW(self.parameters(), lr=self.lr,weight_decay=1e-5)
    scheduler = ReduceLROnPlateau(opt,factor=0.5,patience=10,mode="min",verbose=True)
    return {'optimizer': opt,'lr_scheduler':scheduler,'monitor':'val/loss'}
    #return opt

  def train_dataloader(self):
    ds = Dataset(data=train_files, transform=train_transforms)
    loader=DataLoader(ds, batch_size=self.bs, shuffle=True,num_workers=self.numworker, collate_fn=list_data_collate)
    return loader

  def training_step(self,batch,batch_idx):
    image,segment=batch["image"], batch["label"]
    out=self(image)
    loss=self.criterion(out,segment)
    outputs = [self.post_pred(i) for i in decollate_batch(out)]
    labels = [self.post_label(i) for i in decollate_batch(segment)]
    dice=dice_coef_multilabel(torch.stack(labels),torch.stack(outputs))
    self.log('train/loss', loss, on_epoch=False,prog_bar=True)
    self.log('train/dice', dice, on_epoch=False,prog_bar=True)
    return loss

  def val_dataloader(self):
    ds = Dataset(data=train_files, transform=val_transforms)
    loader=DataLoader(ds, batch_size=1, shuffle=False,num_workers=self.numworker, collate_fn=list_data_collate)
    return loader
    
  def validation_step(self,batch,batch_idx):
    image,segment=batch["image"], batch["label"]
    roi_size = (64, 64, 64)
    sw_batch_size = 4
    outputs = sliding_window_inference(image, roi_size, sw_batch_size, self.forward)
    loss=self.criterion(outputs,segment)
    outputs = [self.post_pred(i) for i in decollate_batch(outputs)]
    labels = [self.post_label(i) for i in decollate_batch(segment)]
    dice=dice_coef_multilabel(torch.stack(labels),torch.stack(outputs))
    self.log('val/dice', dice, on_epoch=True,prog_bar=True)
    self.log('val/loss', loss, on_epoch=True,prog_bar=True)
    return loss
  

[Nic-Ma]

Hi @talhaanwarch ,

Thanks for your interest here.
I think it may be expected behavior as it failed to converge?
@rijobro @wyli Do you have more training suggestions?

Thanks.

[Breeze-Zero]

Hi, I customized a model for Brats21 segmentation and used DiceCELoss under PyTorch Lighting amp. After training for some time, my detection indicated NAN in softMax calculation, but using the model on MONAI did not have these problems. Do you have any solution? My model already uses monai components for the most part.

[Jamshidhsp]

@talhaanwarch Hi, Could you manage to work it out?
I have the same problem using UNETR on segmentation.

[rijobro]

1e-3 is quite a large learning rate. Try with more like 1e-5 first, if that works, then you can increase slightly.

[rijobro]

Also I can't see AMP being used, which is good for debugging. But just to be safe: when getting NaNs, AMP should be disabled as this can sometimes be the cause.

[rijobro]

Also:

• You have Resized but do not have interpolation to be "linear", "nearest".
• You have both Spacing and Resize. This is will cause your images to be resampled multiple times, increasing resampling errors.
• Whilst debugging, simplify things as much as possible:
  • switch from UNETR to UNet
  • remove the LR scheduler
  • switch the loss from DiceFocal to Dice
  • Remove any random transformations

[matt3o]

tl;dr Add the SignalFillEmpty transforms to remove any NaNs after the input transforms

We have discovered one possible reason for this problem. In our case the NaNs errors only triggered when amp was on. Switching it off made the error disappear and training worked as expected. Nevertheless my mentor @Zrrr1997 dug a bit deeper and found out our input already contained NaNs.
This happens especially for crops at the border where our AutoPET volumes were almost completely 0 or maybe even completely 0.
We are still in the process of figuring out why that happens, we think there is some normalization with a division by zero being done, which then in turn creates the NaNs in the input data for the network (so after the pre-transforms).

The solution: I created a SignalFillEmptyd analouge to SignalFillEmpty, which casts all input NaN values to 0 before passing the input to the network. I will PR the transform in the next days and hope this helps other people as well.
-> The network now trains perfectly with AMP switched on.

As to why this bug only triggers with AMP, more debugging is necessary, so far we don't understand it either.

Some additional information, maybe that also helps: I only saw this error pop up on the A100 and I have no clue why, On Rtx 3090 Ti, A6000 and even H100 the error did not occur (at least regularily). This could be a hardware bug for all I know, at least the occurence of this problem is highly hardware related.

If we gather any more information I will update this post.

Update: SignalFillEmptyd got accepted and should be available in the next MONAI release, for anyone who wants to try it out

[wyli]

perhaps check the tf32 flags as well https://docs.monai.io/en/latest/precision_accelerating.html#settings

[matt3o]

I did set those to false (had tf32 for cuda and cudnn on), made no difference, that was my initial guess too. So still crashes as usual with amp on.

[matt3o]

Btw @wyli I was wondering since a longer time, is there any transform for debugging the data dictionary? So printing the different entries along with some information about the tensors (size in memory, gpu or ram, NaN count / warning)?

[wyli]

yes there's DataStats and DataStatsD

MONAI/monai/transforms/utility/array.py

Line 670 in e5f9337

class DataStats(Transform):