nan training loss

[nourmagde00]

Hi ,
I get a train loss equal nan when I train my model , I tried to solve it for several dayes but it does't work yet,
Could u check my notebook please ?

SEG NetJupyter Notebook.pdf

Thanks In Advance.

[rijobro]

The code is too long for us to reasonably debug a PDF file. Try to make a minimum working example based off our tutorials, and then increase the complexity until you end up solving the research question that you are currently interested in.

One bug I noticed -- your loss function uses sigmoid=True, but I think this should be softmax=True since you have multi-channel output.

[nourmagde00]

Hi @rijobro I'd like to thank u for your replying :
this is my Transforms :
train_transform = Compose(
[
# load 4 Nifti images and stack them together
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys="image"),
ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(
keys=["image", "label"],
pixdim=(1.0, 1.0, 1.0),
mode=("bilinear", "nearest"),
),
RandSpatialCropd(keys=["image", "label"], roi_size=[224, 224, 144], random_size=False),
CropForegroundd(keys=["image", "label"], source_key="image"),
Resized (keys=['image', 'label'],spatial_size=(128, 128, 128),mode=['area', 'nearest']),
SpatialPadd(keys=["image", "label"], spatial_size=[128,128,128]),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
RandScaleIntensityd(keys="image", factors=0.1, prob=1.0),
RandShiftIntensityd(keys="image", offsets=0.1, prob=1.0),
EnsureTyped(keys=["image", "label"]),
]
)
val_transform = Compose(
[
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys="image"),
ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(
keys=["image", "label"],
pixdim=(1.0, 1.0, 1.0),
mode=("bilinear", "nearest"),
),
CropForegroundd(keys=["image", "label"], source_key="image"),
Resized (keys=['image', 'label'],spatial_size=(128, 128, 128),mode=['area', 'nearest']),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
EnsureTyped(keys=["image", "label"]),
]
)

and this is my model after change sigmoid with softmax :

max_epochs = 100
val_interval = 1
VAL_AMP = True

standard PyTorch program style: create SegResNet, DiceLoss and Adam optimizer

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = SegResNet(
blocks_down=[1, 2, 2, 4],
blocks_up=[1, 1, 1],
init_filters=16,
in_channels=4,
out_channels=3,
dropout_prob=0.2,
).to(device)
loss_function = DiceLoss(smooth_nr=0, smooth_dr=1e-5, squared_pred=True, to_onehot_y=False, softmax=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-4, weight_decay=1e-5)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_epochs)

dice_metric = DiceMetric(include_background=True, reduction="mean")
dice_metric_batch = DiceMetric(include_background=True, reduction="mean_batch")

post_trans = Compose(
[EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)]
)

define inference method

def inference(input):

def _compute(input):
    return sliding_window_inference(
        inputs=input,
        roi_size=(128, 128, 128),
        sw_batch_size=1,
        predictor=model,
        overlap=0.5,
    )

if VAL_AMP:
    with torch.cuda.amp.autocast():
        return _compute(input)
else:
    return _compute(input)

use amp to accelerate training

scaler = torch.cuda.amp.GradScaler()

enable cuDNN benchmark

torch.backends.cudnn.benchmark = True

[dongyang0122]

@nourmagde00 it is a known issue that AMP does not support very well with 3D operations. Sometimes turning on AMP would return NaN loss values, especially for SegResNet. You can either turn off AMP, or switch to a different segmentation network.

[etienne87]

really?? can you point to a source/ link for this issue? amp is critical for me as i am training on very big volumes