Problem training SwinUNETR segmentation

[nicktasios]

I mostly copied the code from here to train a SwinUNETR for segmentation of a single label. You can see the complete code here.
While training using this code, I get very weird loss progression:

Do you have any suggestions as for why this could be happening? I already tried training without AMP.

[tangy5]

Hi @nicktasios , thanks for the question.
From the loss curve, the model is not learning/converging.
I saw the code, some parameters are very different from training CT images, not sure if it's the data processing issue.

Some suggestions to check:

1. The initial LR is set to 1e-3, which is a little bit high for typical segmentation tasks. Could try 1e-4
   ` parser.add_argument("--optim_lr", default=1e-3, type=float, help="optimization learning rate")
2. output channel is 1, but default DiceCELoss included background, there might causing errors. The model is not optimizing DiceCELoss correctly. Might should be 2?
3. The scale range and spacing are very different, not sure if this will cause problem.

Thanks.
`

[nicktasios]

@tangy5 Thank you very much for your reply. I tried different learning rates, so 1 didn't seem to be the problem, and I know that the scaling works fine on different frameworks. Possibility 2 was interesting so I tried it oud, and indeed it seems to be the problem.

Now I am a bit confused on how Monai works. Is it possible to setup a segmentation pipeline using a single channel where 0s are the background and 1s are the foreground? If yes, what do I need to modify in the code above to make it work properly?

[tangy5]

Hi @nicktasios , good to know the further steps.
About the segmentation pipeline, the answer is Yes. You could do background as 0 and 1 for foreground.

I saw you were coding with the research contribution repo, there is a tutorial as well in the "Tutorials" repo, please refer to this for details:
https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/swin_unetr_btcv_segmentation_3d.ipynb
You could see in this tutorial, there are 14 channels, 0 is the background, and 1-13 for different foreground classes.
If you only has background and one foreground class, set output channels to 2 will be good enough.

Thanks

[nicktasios]

Hi @tangy5 . I'd like to use a single channel as output, not 2. So background has the value of 0 and foreground has the value of 1.

[tangy5]

Yes, you could use sigmoid option for your loss function, set output channel to 1. Then set a threshold when generating the mask, such as 0.5. >0.5 will be the foreground, otherwise background.

Sample Sigmoid usecase: https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/swin_unetr_brats21_segmentation_3d.ipynb