Using SwinUNETR for Binary lesion Segmentation

[burhr2]

Hello, has anyone tried to train SwinUNETR from scratch for binary segmentation? I have tried from my side and after more than 1500 epochs I still ended with the model predicting background class only. The validation loss seemed to have been dropping but with strange behaviour i.e. it's kinda stuck oscillating in a local minimal then sudden sharp drop then stuck again. First time working with Transformers hence I may have missed some important info to consider for training. I used a default setup as suggested in this tutorial and only changed the in_channels=1 and out_channels=2 as I am working on binary segmentation. In advance, thanks for the suggestions.

unet = SwinUNETR(
            img_size=(128, 128, 128),
            in_channels=1,
            out_channels=2,
            feature_size=48,
            drop_rate=0.0,
            attn_drop_rate=0.0,
            dropout_path_rate=0.0,
            use_checkpoint=False,
        )
model = Model(unet, loss_function, config)

context:
We are working on MS lesion segmentation from Spinal cord MRI. The most noticeable challenge is the imbalance btn foreground and background voxels. We have existing Unet solutions which work well but we thought to try transformer architectures as they are reported to outperform conv-only UNets. We are using Torchio + PyTorch lighting frameworks. The dataset is small (62 training subjects) but using patching strategy to improve the sample size. More details on our task can be on our recent published work.

references

1. tutorial link: https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/swin_unetr_brats21_segmentation_3d.ipynb
2. https://inria.hal.science/hal-03865212/

I have posted a similar question on Q&A Discord channels. Thanks

[KumoLiu]

Hi @burhr2, thanks for your interest here.
I noticed that in the tutorial, it hard-coded the threshold to 0.5 in post_pred = AsDiscrete(argmax=False, threshold=0.5), I would suggest you change all the activation functions to softmax in this tutorial and set argmax=True.
BTW, you mentioned that your data is imbalanced, I would also suggest you use RandCropByPosNegLabeld to crop the patches based on the Pos Neg ratio.

Hope it helps, thanks!

[burhr2]

Hi @KumoLiu thank you for the suggestions. I have all the suggestions in place for our current torchio pipeline (maybe to increase foreground label sampling as we currently have:

sampler = tio.data.LabelSampler(patch_size, 'segmentation', label_probabilities={0: 0.5, 1: 0.5})

• I have tested the pipeline with the MONAI UNet and the network is training as expected i.e. train_loss and val_loss is decreasing etc.
• I also observed that when I trained the SwinUNETR for a few epochs (20), I can see the model is somehow trying to segment the spinal cord. Leaving it to train for longer epochs that when it collapses and only predicts the background class.
• I am doing fine-tuning based on this tutorial and see if that could help. But looking at the current train_val_loss does not seem to work well either. Maybe I should try to preprocess my data in a different way etc or I am missing some specific configurations for my task.

Update:

• I have fine-tuned the SwinUNETR model based on this tutorial and the model can now detect lesions, however few points to note:

• the number of false positives is more than 5X higher than our current baseline. This may be due to the very small lesions that need to be segmented which are very different from most of the current medical imaging data (including MS lesions in the brain).

• Training from scratch was not a good idea since we are working on a smaller dataset and that may be the reason why the model ended up predicting the background only.

• Reducing model complexity (size) and the way we preprocess the data may help but will welcome other suggestions and possibly collaboration for the specific task.