Strange training/validation curves for Swin transformer SSL pre-training

[fengling0410]

Hi all, I'm trying to reproduce the pre-trained Swin UNETR and would like to apply it to internal CT scans in my institution. I copy exactly the same code from this address except the dataset directory part of data_utils.py. The datasets I used are LUNA_16, LIDC, TCIA_colon, and TCIA_COVID19, as stated in the SSL Swin transformer paper. I run the code using single gpu with the command python main.py --use_checkpoint --batch_size=1 --num_steps=10000 --lrdecay --eval_num=100 --logdir='./runs/test' --lr=4e-4 roi_x=96 --roi_y=96 --roi_z=96 --lr_schedule='warmup_cosine'. However, the training plot seems quite strange:
The total loss and recon loss doesn't seem to decrease as training goes.

The validation image is also strange and it seems like the model is not learning useful features:

May I ask for some advices regarding this? I'm still new to CT scan segmentation and a little confused about situation. If you spotted anything wrong please let me know. I would appreciate any help from the community :)

[tangy5]

Hi Fengling410, I'm not sure with this issue. But when batch size is set to 1, the contrastive learning will cause some problem that lead to unstable overall training. You could try increase the batch size and see whether it will stablize the training. Thanks.

[fengling0410]

Hi @tangy5, thanks for your reply. Since I'm using Tesla T4, which only works for batchsize=1, I cannot really increase the batch size. But I made use of 4 GPUs for pre-training following your advice and the Swin Transformer works! I guess the reason is that Swin Transformer needs much larger batch size and training iterations than the common ViT structure.