UNETR validation metric stuck on 0.5

[carlpe]

Hi,

I have the following code for training a UNETR 3D segmentation using a single class:

from monai.utils import first, set_determinism
from monai.transforms import (
    AsDiscrete,
    EnsureChannelFirstd,
    Compose,
    CropForegroundd,
    LoadImaged,
    Orientationd,
    RandFlipd,
    RandCropByPosNegLabeld,
    RandShiftIntensityd,
    ScaleIntensityRanged,
    Spacingd,
    RandRotate90d,
)
from monai.losses.dice import DiceCELoss
from monai.networks.nets import UNETR
from tqdm import tqdm
from monai.networks.layers import Norm
from monai.metrics import DiceMetric
from monai.losses import DiceLoss
from monai.inferers import sliding_window_inference
from monai.data import CacheDataset, DataLoader, Dataset, decollate_batch
import torch
import matplotlib.pyplot as plt
import os
import glob
from monai.optimizers import Novograd
from monai.metrics import DiceMetric
from monai.utils import set_determinism

CUDA_LAUNCH_BLOCKING = "1"
data_root = "/home/"
path_to_target = "/home/"

train_images = sorted(glob.glob(os.path.join(path_to_target, "ct", "*.nii*")))
train_labels = sorted(glob.glob(os.path.join(path_to_target, "gt", "*.nii*")))

data_dicts = [
    {"image": image_name, "label": label_name}
    for image_name, label_name in zip(train_images, train_labels)
]
train_files, val_files = data_dicts[:90], data_dicts[90:93]
print(len(train_files))
print(len(val_files))
print("------------")
print("last img:" + train_images[-1])
print("last seg" + train_labels[-1])
print("------------")


train_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(1, 1, 1),
            mode=("bilinear", "nearest"),
        ),
        CropForegroundd(keys=["image", "label"], source_key="image"),
        RandCropByPosNegLabeld(
            keys=["image", "label"],
            label_key="label",
            spatial_size=(96, 96, 96),
            pos=1,
            neg=1,
            num_samples=4,
            image_key="image",
            image_threshold=0,
        ),
        RandFlipd(
            keys=["image", "label"],
            spatial_axis=[0],
            prob=0.10,
        ),
        RandFlipd(
            keys=["image", "label"],
            spatial_axis=[1],
            prob=0.10,
        ),
        RandFlipd(
            keys=["image", "label"],
            spatial_axis=[2],
            prob=0.10,
        ),
        RandRotate90d(
            keys=["image", "label"],
            prob=0.10,
            max_k=3,
        ),
        RandShiftIntensityd(
            keys=["image"],
            offsets=0.10,
            prob=0.50,
        ),
    ]
)
val_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(1, 1, 1),
            mode=("bilinear", "nearest"),
        ),
        CropForegroundd(keys=["image", "label"], source_key="image"),
    ]
)

check_ds = Dataset(data=train_files, transform=train_transforms)
check_loader = DataLoader(check_ds, batch_size=1)
check_data = first(check_loader)
image, label = (check_data["image"][0][0], check_data["label"][0][0])
print(f"image shape: {image.shape}, label shape: {label.shape}")

train_ds = CacheDataset(
    data=train_files, transform=train_transforms, cache_rate=1.0, num_workers=4
)
train_loader = DataLoader(train_ds, batch_size=3, shuffle=True, num_workers=1)
val_ds = CacheDataset(
    data=val_files, transform=val_transforms, cache_rate=1.0, num_workers=0
)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=0)

device = torch.device("cuda:0")

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = UNETR(
    in_channels=1,
    out_channels=2,
    img_size=(96, 96, 96),
    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,
).to(device)

loss_function = DiceCELoss(to_onehot_y=True, softmax=True)
model.load_state_dict(torch.load("mect_unetr_1class.pth"))
torch.backends.cudnn.benchmark = True
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)


def validation(epoch_iterator_val):
    model.eval()
    with torch.no_grad():
        for batch in epoch_iterator_val:
            val_inputs, val_labels = (batch["image"].cuda(), batch["label"].cuda())
            val_outputs = sliding_window_inference(val_inputs, (96, 96, 96), 4, model)
            val_labels_list = decollate_batch(val_labels)
            val_labels_convert = [
                post_label(val_label_tensor) for val_label_tensor in val_labels_list
            ]
            val_outputs_list = decollate_batch(val_outputs)
            val_output_convert = [
                post_pred(val_pred_tensor) for val_pred_tensor in val_outputs_list
            ]
            dice_metric(y_pred=val_output_convert, y=val_labels_convert)
            epoch_iterator_val.set_description(
                "Validate (%d / %d Steps)" % (global_step, 10.0)
            )
        mean_dice_val = dice_metric.aggregate().item()
        dice_metric.reset()
    return mean_dice_val


def train(global_step, train_loader, dice_val_best, global_step_best):
    model.train()
    epoch_loss = 0
    step = 0
    epoch_iterator = tqdm(
        train_loader, desc="Training (X / X Steps) (loss=X.X)", dynamic_ncols=True
    )
    for step, batch in enumerate(epoch_iterator):
        step += 1
        x, y = (batch["image"].cuda(), batch["label"].cuda())
        logit_map = model(x)
        loss = loss_function(logit_map, y)
        loss.backward()
        epoch_loss += loss.item()
        optimizer.step()
        optimizer.zero_grad()
        epoch_iterator.set_description(
            "Training (%d / %d Steps) (loss=%2.5f)"
            % (global_step, max_iterations, loss)
        )
        if (
            global_step % eval_num == 0 and global_step != 0
        ) or global_step == max_iterations:
            epoch_iterator_val = tqdm(
                val_loader, desc="Validate (X / X Steps) (dice=X.X)", dynamic_ncols=True
            )
            dice_val = validation(epoch_iterator_val)
            epoch_loss /= step
            epoch_loss_values.append(epoch_loss)
            metric_values.append(dice_val)
            if dice_val > dice_val_best:
                dice_val_best = dice_val
                global_step_best = global_step
                torch.save(
                    model.state_dict(), os.path.join(data_root, "mect_unetr_1class.pth")
                )
                print(
                    f"Model Was Saved ! Current Best Avg. Dice: {round(dice_val_best, 3)} Current Avg. Dice: {round(dice_val, 3)}"
                )
            else:
                print(
                    f"Model Was Not Saved ! Current Best Avg. Dice: {round(dice_val_best, 3)} Current Avg. Dice: {round(dice_val, 3)}"
                )

        global_step += 1
    return global_step, dice_val_best, global_step_best


max_iterations = 250000
eval_num = 1000
post_label = AsDiscrete(to_onehot=2)
post_pred = AsDiscrete(argmax=True, to_onehot=2)
dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
global_step = 0
dice_val_best = 0.0
global_step_best = 0
epoch_loss_values = []
metric_values = []
while global_step < max_iterations:
    global_step, dice_val_best, global_step_best = train(
        global_step, train_loader, dice_val_best, global_step_best
    )
model.load_state_dict(torch.load(os.path.join(data_root, "mect_unetr_1class.pth")))

print(
    f"train completed, best_metric: {dice_val_best:.4f} "
    f"at iteration: {global_step_best}"
)

I was able to train it successfully (DSC=0.80) using two classes. Now I have merged all my classes into a single class and during training, the dice metric is stuck to 0.5 and it will not increase.

I suppose the issue has to do with post_pred, post_label. (onehot or something there) when going from multi-class to single-class.

Please let me know what I am doing wrong and how to fix this.

[tangy5]

The code looks normal, I feel the post label and pred are correct:

post_label = AsDiscrete(to_onehot=2)
post_pred = AsDiscrete(argmax=True, to_onehot=2)
dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)

I guess there is one class Dice is 1 and the other one is all 0, resulting Dice=0.5 all the time. Can you confirm the data and labels are correct after merging classes?

[carlpe]

Thanks @KumoLiu

I will try our suggestions and check the data once more. I also think that it will not be better with one merged class, but we want to try this experiment

[carlpe]

I found the error.
Of course, a mistake in the dataset folder selection when merging of the classes together. It has been fixed now and training.

Thank you again

[carlpe]

Can any of these hyper parameters be tuned in order to look for smaller structures, such as thin vessels?

model = UNETR(
    in_channels=1,
    out_channels=2,
    img_size=(96, 96, 96),
    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,
).to(device)

Glad for suggestions and insights here

[tangy5]

Hi @carlpe , if the structures are small, which means the segmentation is more challenging and harder task, you can adjust the "feature_size", "num_heads" which controls the scale of the model, for example, feature_size=32, or even 48, num_head=24, will lead to a larger models, and more powerful. But this will take more GPU memory. You could try tune these two parameters and fit for the GPU usage.

[carlpe]

Ok, I will try it 😊 thank you

[carlpe]

"It is strongly recommended to change include_background to False since there is no need to count background." -- For this to work, do I have to change some of the parameters like to_onehot to be one less than before, since the background is not counted, or do I leave it like it is?

By the way, in this experiment, using two classes led to better results than using one class (DSC 0.83 vs 0.68). Increasing the feature_size to 32 and num_heads to 24 did produce a better result in both cases.

[KumoLiu]

Hi @carlpe, when changing include_background to False, you don't need to do anything else. It will remove the background automatically.
https://github.com/Project-MONAI/MONAI/blob/6cb4ced0ea50b427a4065b4516be61b7a073abce/monai/metrics/meandice.py#L131-L132

Thanks!

[carlpe]

Ok, thank you @KumoLiu