Swin UNETR reaches memory limit in Tesla T4

[fengling0410]

Hi all, I'm training Swin UNETR model using MONAI package these days but the training memory reaches limit when I was doing distributed training on four of 16GB Tesla T4 GPUs with a batch size of 1 and swin batch size of 1. Then I changed the Swin UNETR model to UNETR with other settings kept all the same, I found that UNETR model works and it only occupies half of the memory required by Swin UNETR. However, as far as I understand, the UNETR model (no.params: 92784046) is much larger than Swin UNETR (no.params: 62187296), and the computation is less complicated in Swin UNETR since it has linear complexity. I'm wondering if there was anything wrong with my implementation, or the Swin UNETR naturally occupies much larger memory during training compared with UNETR. May I ask for some advices regarding this? Thank you in advance :)

[tangy5]

Hi @fengling0410 , if the image input patch size is 96x96x96 or 128x128x128, it should work on 16G GPUs with batch size of 1. The swin transformer is larger than vanilla vision transformer. So the SwinUNETR used 8 layers in total instead of 12 as the encoder. To our observation, Swin transformer also has weigh larger GFLOPS, which increased the complexity.
From your envronment, the Swin UNETR should work with at least batch size of 1. Can you confirm if the input image patch size is appropriate? Also the validation may cause OOM as well, if the original data dimension is large, slidingWindowInference will cache the entire data in the GPU, you could choose to cache the data to CPU if OOM caused by validation. Hope your problem can be resolved soon. thanks.