not good results with UNETR

[elahe-badali]

Hello @ahatamiz, I hope you well.
I am using your model UNETR on a dataset of 3d MRI images to extract a binary mask. I have 202 images and with some augmentation i increase it to about 430 image.
I have mean dice 28% on validation part in about 30 epoch and after this network start to overfit. I try some tricks but no lock.
I want to know your comments. Thank you.

[ahatamiz]

hi @elahe-badali

Do you have any baselines for other model's performance such as UNET on this dataset ? I think each case requires its own tuning and hyper-parameter optimization, but from your descriptions, it maybe that the dataset is quite small and needs additional considerations to avoid overfitting.

One recommendation is to increase weight decay (i.e. --decay in here to larger values. Another recommendations is to add more data augmentation transforms from MONAI ( done on the fly ). Lastly, you may want to use a lighter UNETR model by specifying (--feature_size to be 16 or 8).

Thanks,

[elahe-badali]

@ahatamiz
I have results of a lighter model from my own based on LSTM. That model's results is quite good but it has the same mean dice (surprisingly) .
my model have about 130M parameters, however UNETR have about 500M.
I use the unofficial version of your code in https://github.com/tamasino52/UNETR
I change the last layers based on my datasets needs.

UNETR results:
left one is main mask and right is prediction:

My LSTM-based model results:

[ahatamiz]

Hi @elahe-badali

From these results, it does not seem that UNETR has properly converged. In addition, I would strongly recommend using the official implementation provided in MONAI along with other functionalities (e.g. transforms, loss functions etc.). A good place to start creating your repository based on MONAI is the UNETR research contributions repository:

https://github.com/Project-MONAI/research-contributions/tree/master/UNETR/BTCV

Thanks