Model loss decreases but validation DICE is always 0

[kristianlindholm]

Hi

First of all, I'm pretty new to this so please let me know if there is anything I can provide in a better way.

I'm trying to train a segmentation model on whole-body PET/CT images. I've tried to implement the standard U-net model for this. The output should be just a simple background=0 and tumor=1 label. The problem I'm facing is that even though the training loss is declining, my validation dice score is just 0, and I can't for the love of god figure out what I'm doing wrong.

This is my entire script:

#!/usr/bin/env python3

print("Hello World!")

import numpy as np
from monai.utils import first, set_determinism
from monai.transforms import (
    AsDiscrete,
    AddChanneld,
    Compose,
    ConcatItemsd,
    CropForegroundd,
    LoadImaged,
    Orientationd,
    RandCropByPosNegLabeld,
    ScaleIntensityRanged,
    Spacingd,
    ToTensord,
    AsChannelFirstD,
    RandFlipd,
    NormalizeIntensityd,
    RandScaleIntensityd,
    RandShiftIntensityd,
    CropForegroundd,
    Rand2DElastic,
    RandAffined,
    )
from monai.networks.nets import UNet
from monai.networks.layers import Norm
from monai.metrics import compute_meandice
from monai.losses import DiceLoss, TverskyLoss, DiceCELoss
from monai.inferers import sliding_window_inference
from monai.data import SmartCacheDataset, CacheDataset, DataLoader, Dataset
from monai.config import print_config
from monai.apps import download_and_extract
from monai.optimizers import Novograd
import torch
torch.multiprocessing.set_sharing_strategy('file_system')
import sys
import matplotlib
import matplotlib.pyplot as plt
import tempfile
import shutil
import os
import glob
from glob import glob

import warnings
warnings.filterwarnings("ignore")  # remove some scikit-image warnings
print_config()

root_dir = "//homes/lindholm/monai/data/nets20"
CT_images = sorted(glob("//homes/lindholm/monai/data/nets20/images/*0000.nii.gz"))
PET_images  = sorted(glob("//homes/lindholm/monai/data/nets20/images/*0001.nii.gz"))
train_labels  = sorted(glob("//homes/lindholm/monai/data/nets20/labels/*.nii.gz"))
  
data_dict = [
    {"pet": PET_images, "ct": CT_images, "label": label_name}
    for PET_images, CT_images, label_name in zip(PET_images, CT_images, train_labels)
]

print(data_dict[0])

train_files = data_dict[:-9]
val_files = data_dict[-9:]

print("Create transformers")
train_transforms = Compose(
    [
        LoadImaged(keys=["pet", "ct", "label"]),
        AddChanneld(keys=["pet", "ct", "label"]),
        NormalizeIntensityd(keys="pet", nonzero=True), # Normalize intensity
        NormalizeIntensityd(keys="ct", nonzero=True), # Normalize intensity
        ConcatItemsd(keys=["pet", "ct"], name="image", dim=0),
        CropForegroundd(keys=["image", "label"], source_key="image"),
        RandCropByPosNegLabeld(
            keys=["image", "label"],
            label_key="label",
            spatial_size=(256,256,256),
            pos=5, # I 1 ud af 1+1 vil crop have en label voxel som centrum
            neg=1,
            num_samples=3,
            image_key="image",
            image_threshold=0,
        ),
        
        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0), # Random flip image on axis
        RandAffined(
            keys=['image', 'label'],
            mode=('bilinear', 'nearest'),
            prob=0.2,
            shear_range=(0, 0, 0.1),
            rotate_range=(0, 0, np.pi/16),
            scale_range=(0.1, 0.1, 0.1)),
        ToTensord(keys=["image", "label"]),
    ]
)
val_transforms = Compose(
    [
        LoadImaged(keys=["pet", "ct", "label"]),
        AddChanneld(keys=["pet", "ct", "label"]),
        NormalizeIntensityd(keys="pet", nonzero=True),
        NormalizeIntensityd(keys="ct", nonzero=True),
        ConcatItemsd(keys=["pet", "ct"], name="image",dim=0),
        
        ToTensord(keys=["image", "label"]),
    ]   
)

print("Define dataset loaders")
train_ds = Dataset(data=train_files, transform=train_transforms)
train_loader = DataLoader(train_ds, batch_size=1, shuffle=True, num_workers=4)

val_ds = Dataset(data=val_files, transform=val_transforms)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=2)

print("Create Model, Loss, Optimizer")
# standard PyTorch program style: create UNet, DiceLoss and Adam optimizer
device = torch.device("cuda:0")
model = UNet(
    dimensions=3,
    in_channels=2,
    out_channels=1,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
).to(device)

squared_pred=True, reduction='mean', batch=False)
loss_function = DiceLoss(include_background=False, to_onehot_y=True, softmax=True)

learning_rate = 1e-4
optimizer = torch.optim.Adam(model.parameters(), learning_rate)
scaler = torch.cuda.amp.GradScaler() #Scaled gradient 

print("Execute a typical PyTorch training process")
max_epochs = 10000
val_interval = 2
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = []
metric_values = []
post_pred = AsDiscrete(argmax=True, to_onehot=True, n_classes=2) # argmax turns values into discretes
post_label = AsDiscrete(to_onehot=True, n_classes=2)

for epoch in range(max_epochs):
    print("-" * 10)
    print(f"epoch {epoch + 1}/{max_epochs}")
    model.train() # Tells the model that it's being trained and not used for inference, model.eval()
    epoch_loss = 0
    step = 0
    
    for batch_data in train_loader:
        step += 1
        inputs, labels = (
            batch_data["image"].to(device),
            batch_data["label"].to(device),
        )
        optimizer.zero_grad()

        with torch.cuda.amp.autocast(): # (automated mixed precision)
            outputs = model(inputs)
            loss = loss_function(outputs, labels)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        
        epoch_loss += loss.item()
        print(
            f"{step}/{len(train_ds) // train_loader.batch_size}, "
            f"train_loss: {loss.item():.4f}")
       
    epoch_loss /= step
    epoch_loss_values.append(epoch_loss)
    print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")

    if (epoch + 1) % val_interval == 0:
        model.eval()
        with torch.no_grad():
            metric_sum = 0.0
            metric_count = 0
            for val_data in val_loader:
                val_inputs, val_labels = (
                    val_data["image"].to(device),
                    val_data["label"].to(device),
                )
                roi_size = (256, 256,256)
                sw_batch_size = 4
                
                with torch.cuda.amp.autocast():
                    val_outputs = sliding_window_inference(
                        val_inputs, roi_size, sw_batch_size, model, overlap=0.25,
                    )

                val_outputs = post_pred(val_outputs)
                val_labels = post_label(val_labels)
                value = compute_meandice(
                    y_pred=val_outputs,
                    y=val_labels,
                    include_background=False,
                )
                metric_count += len(value)
                metric_sum += value.sum().item()
                
            metric = metric_sum / metric_count
            metric_values.append(metric)
            if metric > best_metric:
                best_metric = metric
                best_metric_epoch = epoch + 1
                torch.save(model.state_dict(), os.path.join(
                    root_dir, "best_metric_model_unet.pth"))
                print("saved new best metric model")
            print(
                f"current epoch: {epoch + 1} current mean dice: {metric:.10f}"
                f"\nbest mean dice: {best_metric:.10f} "
                f"at epoch: {best_metric_epoch}"
            )
            
    
    if (epoch + 1) % 3 == 0:
        print("Plot the loss and metric")
        
        plt.figure("train", (12, 6))
        plt.subplot(1, 2, 1)
        plt.title("Epoch Average Loss")
        x = [i + 1 for i in range(len(epoch_loss_values))]
        y = epoch_loss_values
        plt.xlabel("epoch")
        plt.plot(x, y)
        
        plt.subplot(1, 2, 2)
        plt.title("Val Mean Dice")
        x = [val_interval * (i + 1) for i in range(len(metric_values))]
        y = metric_values
        plt.xlabel("epoch")
        plt.plot(x, y)
        
        plt.show()
        fname = "//homes/lindholm/monai/data/nets20/metrics_unet.png"
        plt.savefig(fname, dpi=300, facecolor='w', edgecolor='w',
                    format='png', transparent=False, pad_inches=0.1,
        )
    
        torch.save(model.state_dict(), os.path.join(
                    root_dir, "latest_metric_model_unet.pth"))
            
print(
    f"train completed, best_metric: {best_metric:.4f} "
    f"at epoch: {best_metric_epoch}")  

[Nic-Ma]

Hi @ericspod and @wyli ,

Could you please help take a look at this question? I think you guys have much experience in model training.

Thanks.

[ericspod]

Typically for DiceLoss or compute_meandice with a binary label you would set include_background to True since we want to compare the foreground against background, and False only if you have more than one class of label. With multiclass segmentation we want to compare the different labels and ignore the background since it will be the whole image minus any labelled values, and so may dominate the loss calculation if the labels are all relatively small.

Also for DiceLoss you want to use sigmoid rather than softmax, this usage here I think is what you want, and DiceMetric might work better for you as well.

[kristianlindholm]

Thank you very much for the response.

First of all, could you maybe explain what is the difference between sigmoid and softmax in this setting? What difference does it make?

Changing the things you propose made somewhat a difference. The loss function is still going down and the validation Dice is still stuck. The value of the dice score is however at 0.5 now.

[Sere1nz]

@ericspod @kristianlindholm same problem here, I'm using swinunetr to train my data. Train loss goes down but validation metric stucks at 0.5. Did you fix it by changing which settings?
Code snippet:

dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
post_pred = Compose([AsDiscrete(argmax=True, to_onehot=2)])
post_label = Compose([AsDiscrete(to_onehot=2)])
loss_function = DiceLoss(to_onehot_y=True, softmax=True)

  val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]
  val_labels = [post_label(i) for i in decollate_batch(val_labels)]
  # compute metric for current iteration
  dice_metric(y_pred=val_outputs, y=val_labels)

This setting works for Unet training, but in swinunetr the val dice stucks at 0.5.

[ericspod]

The idea with applying sigmoid in the binary case is that we want to convert the logits to something as close to a binary segmentation as possible. Any negative values become background, positive become the segmentation. Softmax is applied in the channel dimension so that a particular pixel/voxel will have the largest value in the channel for the most likely category, applying argmax to this then gives you the segmentation image.

If following our tutorials in the notebooks so that you have a similar setup with the loss isn't helping perhaps there's something wrong with the data you're feeding the metric. I would look at a validation batch with matplotlib to make sure the images are sensible since a score of exactly 0.5 looks to me like something's misconfigured with the validation data pipeline or something else data related.

[kristianlindholm]

Thanks for your response! I'm pretty sure the data is alright. I'm using it for another model with great succes.

I can see that in your 3D segmentation of the spleen, the output no. of channels of the Unet is set to 2. I'm setting mine to 1, as I only want 1 channel in the label. Could you maybe explain why 2 output channels are chosen in the notebook?

I'm thinking the issue has something to do with my model output / post_processing / loss calculation...

[ericspod]

That actually might be a mistake in the notebook. With binary segmentation it does work to treat fore and background as separate labels, I just don't think it's optimal.

[kristianlindholm]

alrigth, well I've now tried to set the number of output channels to two, and this way I can actually get my model to learn and get a somewhat proper output. I'm however still in doubt if I'm doing the right things.

I've changed these lines of code

model = UNet(
    dimensions=3,
    in_channels=2,
    out_channels=2,

loss_function = DiceLoss(include_background=True, to_onehot_y=True, softmax=True )

post_pred = Compose([Activations(sigmoid=True),
                     AsDiscrete(argmax=True, to_onehot=True, n_classes=2)]) 
post_label = AsDiscrete(to_onehot=True, n_classes=2)
dice_metric = DiceMetric(include_background=True, reduction="mean") # For validation

and then in my inference script I'm applying this post prediction transform
post_trans = Compose([Activations(sigmoid=True), AsDiscrete(argmax=True)])

Does this seem somewhat sensible?

[ericspod]

I think this is, though I'm not sure why setting out_channels to 1 and using sigmoid with DiceLoss doesn't work for you, but if this does work then go with it.

[Saba-Heidari]

Hello,
I'm trying to train 3D Unet for CT segmentation with two labels like the above question ( background=0 and tumor=1) and the size of my input images is 256x256x100. I face the same issue, but the above solution is not working. My training loss is declining, but my validation dice score not changing and it's always 1 :(
Could you please help me to figure out the problem?
Thanks

My code:

root_dir = "/Data/images/.nii.gz"))
train_labels = sorted(glob("/Data/labels/.nii.gz"))

data_dict = [
{"image": CT_images, "label": label_name}
for CT_images, label_name in zip( CT_images, train_labels)
]

train_files = data_dict[:-9]
val_files = data_dict[-9:]

train_transforms = Compose(
[
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys=["image", "label"]),
Orientationd(keys=["image", "label"], axcodes="RAS"),
ScaleIntensityRanged(keys=["label"], a_min=0, a_max=256, b_min=0.0, b_max=1.0, clip=False, ),
ScaleIntensityRanged(
keys=["image"], a_min=-57, a_max=164,
b_min=0.0, b_max=1.0, clip=True,
),
CropForegroundd(keys=["image", "label"], source_key="image"),
RandCropByPosNegLabeld(
keys=["image", "label"],
label_key="label",
spatial_size=(128, 128, 64),
pos=1,
neg=1,
num_samples=4,
image_key="image",
image_threshold=0,
),
EnsureTyped(keys=["image", "label"]),
]
)
val_transforms = Compose(
[
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys=["image", "label"]),
ScaleIntensityRanged(keys=["label"], a_min=0, a_max=256, b_min=0.0, b_max=1.0, clip=False, ),
ScaleIntensityRanged(
keys=["image"], a_min=-57, a_max=164,
b_min=0.0, b_max=1.0, clip=True,
),
CropForegroundd(keys=["image", "label"], source_key="image"),
EnsureTyped(keys=["image", "label"]),
]
)

###################################################################################
train_ds = Dataset(data=train_files, transform=train_transforms)
train_loader = DataLoader(train_ds, batch_size=1, shuffle=True, num_workers=4)

val_ds = Dataset(data=val_files, transform=val_transforms)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=2)

device = torch.device("cuda:0")
model = UNet(
dimensions=3,
in_channels=1,
out_channels=1,
channels=(16, 32, 64, 128, 256),
strides=(2, 2, 2, 2),
num_res_units=2,
norm=Norm.BATCH,
).to(device)

loss_function = DiceLoss(include_background=True, to_onehot_y=True, softmax=True, squared_pred=True, reduction='mean', batch=False)
learning_rate = 1e-4
optimizer = torch.optim.Adam(model.parameters(), learning_rate)
scaler = torch.cuda.amp.GradScaler() #Scaled gradient

##################################################################

max_epochs = 10000
val_interval = 2
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = []
metric_values = []
post_pred = Compose([Activations(sigmoid=True), AsDiscrete(argmax=True, to_onehot=True, n_classes=2)]) # argmax turns values into discretes
post_label = AsDiscrete(to_onehot=True, n_classes=2)
dice_metric = DiceMetric(include_background=True, reduction="mean")
post_trans = Compose([Activations(sigmoid=True), AsDiscrete(argmax=True)])

for epoch in range(max_epochs):
print("-" * 10)
print(f"epoch {epoch + 1}/{max_epochs}")
model.train() # Tells the model that it's being trained and not used for inference, model.eval()
epoch_loss = 0
step = 0

for batch_data in train_loader:
    step += 1
    inputs, labels = (
        batch_data["image"].to(device),
        batch_data["label"].to(device),
    )
    optimizer.zero_grad()

    with torch.cuda.amp.autocast():  # (automated mixed precision)
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

    epoch_loss += loss.item()
    print(
        f"{step}/{len(train_ds) // train_loader.batch_size}, "
        f"train_loss: {loss.item():.4f}")

epoch_loss /= step
epoch_loss_values.append(epoch_loss)
print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")

if (epoch + 1) % val_interval == 0:
    model.eval()
    with torch.no_grad():
        metric_sum = 0.0
        metric_count = 0
        for val_data in val_loader:
            val_inputs, val_labels = (
                val_data["image"].to(device),
                val_data["label"].to(device),
            )
            roi_size = (256,256,256)

            sw_batch_size = 4

            with torch.cuda.amp.autocast():
                val_outputs = sliding_window_inference(
                    val_inputs, roi_size, sw_batch_size, model, overlap=0.25,
                )
            val_outputs = post_trans(val_outputs)
            val_outputs = post_pred(val_outputs)

            val_labels = post_label(val_labels)
            print(val_outputs.shape)
            print(val_labels.shape)
            # compute metric for current iteration
            dice_metric(y_pred=val_outputs, y=val_labels)

            # aggregate the final mean dice result
        metric = dice_metric.aggregate().item()

        dice_metric.reset()

        metric_values.append(metric)

        if metric > best_metric:
            best_metric = metric
            best_metric_epoch = epoch + 1
            torch.save(model.state_dict(), os.path.join(
                root_dir, "best_metric_model_unet.pth"))
            print("saved new best metric model")
        print(
            f"current epoch: {epoch + 1} current mean dice: {metric:.10f}"
            f"\nbest mean dice: {best_metric:.10f} "
            f"at epoch: {best_metric_epoch}"
        )

if (epoch + 1) % 3 == 0:
    print("Plot the loss and metric")

    plt.figure("train", (12, 6))
    plt.subplot(1, 2, 1)
    plt.title("Epoch Average Loss")
    x = [i + 1 for i in range(len(epoch_loss_values))]
    y = epoch_loss_values
    plt.xlabel("epoch")
    plt.plot(x, y)

    plt.subplot(1, 2, 2)
    plt.title("Val Mean Dice")
    x = [val_interval * (i + 1) for i in range(len(metric_values))]
    y = metric_values
    plt.xlabel("epoch")
    plt.plot(x, y)

    plt.show()
    fname = "/Data/metrics_unet.png"
    plt.savefig(fname, dpi=300, facecolor='w', edgecolor='w',
                format='png', transparent=False, pad_inches=0.1,
                )

    torch.save(model.state_dict(), os.path.join(
        "/Data/latest_metric_model_unet.pth"))

print(
f"train completed, best_metric: {best_metric:.4f} "

[ebrahimebrahim]

Just chiming in to point out that to_onehot_y=True in

loss_function = DiceLoss(include_background=True, to_onehot_y=True, softmax=True, squared_pred=True, reduction='mean', batch=False)

means that loss_function now expects segmentation labels to not be one-hot encoded, but rather to have a single channel with discrete class labels. This might be conflicting with the decision to have

AsDiscreted(keys=["image", "label"],to_onehot=True)

already transforming labels into a one-hot format.

[Saba-Heidari]

Thank you for your response.
I change the loss function (remove to_onehot_y=True,) and add AsDiscreted(keys=["image", "label"],to_onehot=True) to my transformation, but still getting the same error as above.
I tried it without having AsDiscreted(keys=["image", "label"],to_onehot=True) in my transformation and the Dice score result is still 1. Do you have any suggestions?

[ebrahimebrahim]

I would inspect intermediate outputs and sanity-check them to make sure they have the expected format.
You can take an example data item and apply transforms one by one, checking that the format is as you expect.
Some good things to sanity-check:

• tensor shape
• min and max values in tensor
• the set of unique values in the tensor (e.g. using torch.unique)

[ericspod]

We have tutorials on 2D segmentation that are a bit simpler compared to what you have but I think they demonstrate how the loss function is set up in a way that should work for you here. With these arguments the dice will operate in a binary formulation. It also has AsDiscrete(threshold=0.5) in the post transformation but I'm not sure that's what you would want, but do try DiceLoss(sigmoid=True) for your loss and see if that helps.

[Saba-Heidari]

It worked.
Thank you

[TuanDTr]

Hi @ericspod, apologize for direct tagging you in the post. I have a similar problem with binary segmentation. I am training UNETR on 2 segmentation tasks, spleen (MSD Task 09) and lung tumor (MSD Task 06). The validation dice score for the lung tumor task is always around 0.00xx and does not improve along the training (after 40 epochs, each epoch of 250 iterations with batch size of 6). The validation dice score for spleen task, though relatively small, shows sign of improvement across time. I wonder if it is because of the size of the object that makes it difficult for the model to converge. Because the size of the lung tumor is pretty small (the median is like 30x40x18) compared to the spleen. Are you aware of similar problems or any techniques to train under these circumstances?

[ericspod]

The initial idea I get is that the size difference is a problem. I have seen in the past that trying to segment small areas in large images is difficult because you just have so little signal (the ground truth segmentation). For the DiceLoss class this is why we have ignore_background for multiclass segmentation, this causes the function to not consider the background class since it's always the absense of the other classes and would dominate the calculation if used. If I'm getting you right your network is multiclass in that you are trying to segment the spleen and the lung tumours, it's possible that your input to the DiceMetric isn't formulated correctly or that your postprocessing sequence isn't correct for multiclass, ie. the values for AsDiscrete aren't correct. I haven't used UNetR much so I can't say if it's behaviour is different enough to account for these problems either.

[TuanDTr]

Hi @ericspod, thanks for your response. I trained two independent tasks for lung tumor and spleen segmentation, i.e. there are two separate datasets for binary segmentation. At the beginning I started just by segmenting the lung tumors and observed weird behavior in validation dice score. So to make sure the model had been correctly implemented, I trained on the spleen segmentation task (bigger object than the tumor) and observed the low dice score but at least it slowly improved over time. Therefore I presume that it has something to do with the object size.

[ericspod]

There's nothing that immediately comes to mind I can suggest, if it's something wrong with the network or something with your training setup I can say. I would suggest going back to one of the tutorials such as spleen segmentation and substitute the given network for UNetr and see if you see the same behaviour, if you do then it's most likely some property of the network to be concerned about or its interaction with the training setup. @ahatamiz would you have any insights?