Diverging Dice Metric in Validation Phase

[Kiarashdnsh]

Hello,

I'm a novice in AI and image segmentation.
I have a dataset of [256, 256, 256]-sized images and labels. The input channel is 1, and the output channel is 3 (label IDs are 0, 1, and 2, respectively). I'm trying to segment the images using UNet within the MONAI framework.

I have used the UNet pipeline provided in the tutorials for my training phase. Based on the mean dice metric, the network has achieved a value of about 0.68 in a 50-epoch training run with a dataset of 160 training images:

Now for the validation phase, I expect the UNet to reach a dice value similar to the dice value of the training phase. However, I get a diverging behavior of the network which reaches only 0.06. For the validation phase, I used the following code:

if len(devices) > 1:
    model = torch.nn.DataParallel(model, device_ids=devices)

post_label = Compose([AsDiscrete(to_onehot=3)])
model.eval()
with torch.no_grad():
    for val_data in val_loader:
        val_inputs, val_labels = val_data["image"].to(device), val_data["label"].to(device)
        # define sliding window size and batch size for windows inference
        roi_size = (96, 96, 96)
        sw_batch_size = 4
        val_outputs = sliding_window_inference(val_inputs, roi_size, sw_batch_size, model)
        val_outputs = [post_trans(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)
        for val_output_dice in val_outputs:
            saver(val_output_dice)
        IoU_metric(y_pred=val_outputs, y=val_labels)
        for val_output_IoU in val_outputs:
        #    saver(val_output_IoU)
    # aggregate the final mean dice result
    print("evaluation dice metric:", dice_metric.aggregate().item())
    print("evaluation IoU metric:", IoU_metric.aggregate().item())
    # reset the status
    dice_metric.reset()
    IoU_metric.reset()

In the end, I got the following result:

evaluation dice metric: 0.06054682657122612
evaluation IoU metric: 0.19970453202724457

So I thought maybe the network was overfitting to the training data and couldn't segment new data with good efficiency. So in the next training phases, I reduced the epochs to 30, 20, 15, and 10 and while the dice value in each training run was less than 0.67, it would get lowered by reducing the number of epochs. In the end, I'd always get the exact same evaluation dice value in the validation phase:

evaluation dice metric: 0.06054682657122612
The only difference was that the IoU value would change between 0.07 and 0.19 in each validation run.

So now, I was sure that the training phase was not overfitting. Thus, I changed the images used in the validation phase with the ones that the network was actually trained with. The result, again was 0.06054682657122612.

To make sure that nothing was wrong in the validation phase, I added
print("evaluation dice metric:", dice_metric.aggregate().item())
to the validation code:

        ...
        dice_metric(y_pred=val_outputs, y=val_labels)
        for val_output_dice in val_outputs:
            saver(val_output_dice)
            print("evaluation dice metric:", dice_metric.aggregate().item())
        ...

so after each segmentation in the validation phase, I get to see the efficiency of that segmentation (not sure it's the correct way though). Then again, I get the same dice values between 0.052 and 0.062, most of them being 0.060. Here's a bar plot of outputted dice values:

Any help regarding solving this issue would very much be appreciated!

Thanks in advance,
Kiarash

[KumoLiu]

Hi @Kiarashdnsh, thanks for your interest here.
Did you use the same post-processing for the training and validation phase? I would suggest you plot the outputs of the model and the outputs that have been activated.
I hope it helps, thanks!

[Kiarashdnsh]

Hi @KumoLiu, thanks for your reply,

If I get your meaning right, I've used the following lines as the post processing code in my network:

...
post_pred = Compose([AsDiscrete(argmax=True, to_onehot=3)])
post_label = Compose([AsDiscrete(to_onehot=3)])
post_trans = Compose([Activations(sigmoid=True)], [AsDiscrete(threshold=0.5)])
...

In the model block where I defined the model of the network and its parameteres.

[KumoLiu]

What are post_pred and post_trans used for individually?
If you used sigmoid I think you output only one channel, so why could you still use argmax?

[Kiarashdnsh]

Since I'm a novice, I don't know why I have to use post_trans and post_pred lines individually, only that they are used in the training:

...
sw_batch_size = 4
                val_outputs = sliding_window_inference(val_inputs, roi_size, sw_batch_size, model)
                val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]
                val_labels = [post_label(i) for i in decollate_batch(val_labels)]
...

and validation block:

...
for val_data in val_loader:
        val_inputs, val_labels = val_data["image"].to(device), val_data["label"].to(device)
        # define sliding window size and batch size for windows inference
        roi_size = (96, 96, 96)
        sw_batch_size = 4
        val_outputs = sliding_window_inference(val_inputs, roi_size, sw_batch_size, model)
        val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]
        val_labels = [post_label(i) for i in decollate_batch(val_labels)]
...

However, my output has three labels; IDs are from 0 to 2. However, if you think I could improve the efficiency of the netwotrk by changing the parameters to exact defintions, please tell me. As I wrote in the intro, I have a dataset of [256, 256, 256]-sized images and labels. The input channel is 1, and the output channel is 3 (label IDs are 0, 1, and 2, respectively).

[KumoLiu]

Hi @Kiarashdnsh, I think you can replace all post_trans with post_pred.
For more segmentation details, you can also refer to https://github.com/Project-MONAI/tutorials/tree/main/3d_segmentation.

Thanks.

[Kiarashdnsh]

Thanks, I'll do that.
One minor question though:
What is the function of defining roi_size and sw_batches in the training and the validation phases? I feed 256 by 256, by 256 images to the network as its input. By defining roi_size I assume the network train on 96 by 96 by 96 croped image sizes. Doesn't this reduce the efficiency of segmentation? Also, I don't know what does sw_batches do? I tried changing its value to see its impact on segmentation. However I couldn't figure out what it does and how it effects the segmentation.

if (epoch + 1) % val_interval == 0:
        model.eval()
        with torch.no_grad():
            for val_data in val_loader:
                val_inputs, val_labels = (
                    val_data["image"].to(device),
                    val_data["label"].to(device),
                )
                roi_size = (96, 96, 96)
                sw_batch_size = 4
                val_outputs = sliding_window_inference(val_inputs, roi_size, sw_batch_size, model)
                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)
                IoU_metric(y_pred=val_outputs, y=val_labels)