U-Net Breast MRI: Low validation score, decreasing train loss

[Marinos736]

I have a tumour segmentation problem and im using the MONAI framework to load and transform my data as well as create a model. The dataset consists of 99 images (512x512x116) with their corresponding binary masks showing benign and malignant tumours. The images are shuffled and split into train and validation sets at 80% split.

Following the Spleen 3D tutorial the code is pretty standard monai-code with the following differences:

Im using minimal transformations for sanity check plus some augmentation

train_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        AddChanneld(keys=["image", "label"]),
        ScaleIntensity(keys=["image"], minv=0.0, maxv=1.0),
        SpatialCropd(keys=["image","label"], roi_start=(0,0,0),roi_end=(512,512,96)
        RandCropByPosNegLabeld(
            keys=["image", "label"],
            label_key="label",
            spatial_size=(256, 256, 96),
            neg=0,
            num_samples=4,
            image_key="image",
            image_threshold=0,
        ),
        ToTensor(keys=["image", "label"]),
    ]
)


val_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        AddChanneld(keys=["image", "label"]),
        ScaleIntensity(keys=["image"], minv=0.0, maxv=1.0),
        SpatialCropd(keys=["image","label"], roi_start=(0,0,0),roi_end=(512,512,96)
        ToTensor(keys=["image", "label"]),
    ]
)

2.ThreadDataLoader instead of Dataloader:

train_loader = ThreadDataLoader(train_ds, batch_size=1, num_workers=0)

3. Loss Function (either GeneralizedDiceLoss or DiceCELoss) and optimizer:

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)
loss_function = GeneralizedDiceLoss(include_background=False, to_onehot_y=True, softmax=True)

4. Post prediction activations:

post_pred = Compose([Activations(softmax=True), AsDiscrete(argmax=True, to_onehot=2)])
While train loss is decreasing steadily, the validation score never surpasses a dice score of ~0.2-0.3 max at over 100 epochs (with training loss under 0.2 at that point).

QUESTIONS

Firstly, is the dataset big enough to train a U-Net from scratch?

Which augmentations would be enough in order to overcome overfitting issues?

The big question is what could be the reason for the low dice score and what can I do?

Also, when im trying to increase batch size i get CUDA error: out of memory so increasing batch size might not be possible for this problem.

I will appreciate any advice on these issues and please do not hesitate to tell me if any additional information is needed.

[KumoLiu]

Hi @Marinos736,

Is the dataset big enough to train a U-Net from scratch?

A dataset of 99 images is quite small for training a deep learning model from scratch, such as U-Net. However, it might still work. The U-Net architecture is specifically designed for semantic segmentation tasks and it's often capable of producing good results even with small datasets, especially in medical imaging.

Which augmentations would be enough to overcome overfitting issues?

Data augmentation can be very useful in preventing overfitting, especially with small datasets. For your medical imaging task, you can consider geometric transformations like rotation, scaling, and flipping. Elastic deformations can also simulate anatomical variations and are commonly used in medical image analysis. Additionally, you can try intensity transformations like random Gaussian noise, random bias field, or gamma adjustments to account for variations in illumination, contrast, and shading.

What could be the reason for the low dice score and what can you do?

The low Dice score could be due to several reasons including small dataset size, model overfitting, insufficient training, improper data preprocessing or augmentation, imbalance between classes, etc.
To improve the Dice score, you might consider:

• Increasing the size of your dataset, if possible. More data can improve the ability of the model to generalize.
• Incorporating regularization techniques to reduce overfitting, such as dropout layers in your network or adding weight decay.
• Ensuring your training process is sufficient. Try training for more epochs, or adjusting the learning rate.
• Tuning your data augmentation strategies to include more variations.
• If there's a class imbalance in your segmentation masks, consider using weighted or focal loss.

Also, when im trying to increase batch size i get CUDA error: out of memory so increasing batch size might not be possible for this problem.

You can use mixed precision training, which halves the memory footprint of your floating-point tensors, and hence allows you to train larger models or use larger batch sizes. Lastly, you could consider using multiple GPUs if you have access to more.

Hope it helps, thanks.

[Marinos736]

Thank you so much for the swift and detailed answer!! I will try the suggestions you mentioned.

[SomeUserName1]

A common way to increase dataset size and decrease memory footprint is to use patch-based data pipelines as provided by torchio: Patch-based pipeline

[Marinos736]

Thank you for your input! I will check it out since dataset size might be a serious issue for my task.

I have another question if someone would be so kind to answer...

For this semantic segmentation task of breast tumours I am a little confused whether I should set out_channels=1 and work with
loss_function = monai.losses.DiceLoss(include_background=False, sigmoid=True)
and something like
post_pred = Compose([Activations(sigmoid=True), AsDiscrete(threshold_values=True)])
at post prediction.

Basically I want a single segmentation mask as output showing the tumours and I believe I need to exclude background from calculations since the tumour area is extremely small when compared to the background.

Is the Spleen 3D tutorial the right fit for this problem? (Since it has out_channels=2)

[Lucas-rbnt]

Hello, if you don't mind me providing some insights about your concerns:

• setting out_channels=1 seems to be the right thing here since there is also one label.
• regarding include_background, the documentation states that setting it to False can help in cases where the tumour masks are small compared to the original image (so your case).
• For your post prediction compose, it basically not interfere with your training as it is only used for metrics computation. In any case, it's the right thing to do, it is possible to set a threshold value in AsDiscrete (see the thresholdargument).
  Why is this threshold value important when calculating metrics? -> Let's assume that the model is also trained on each pixel in the background (include_background=True). The dataset is small, which may affect the model's ability to generalise properly. In addition, with a very small mask area, the model may choose, in order to minimise the cost function (which is what we are asking it to do), to bring all its predictions closer to 0 for each pixel. But it can still provide values around 0.25/0.30 for the pixel values in the tumour mask (pure guess for these). By setting the threshold to 0.5, all pixels are written as predicted to be 0. The model is bad, but by setting this value to 0.2 you get a bloody good model. Hence, you can inspect the distribution of predicted values for your mask before discretization to get a clearer idea of how your model works.