UNETR memory allocation issue

[Pyromonk]

Good day.

Is it normal for UNETR to try to occupy unreasonable amounts of memory? Or am I not using it in the right way somehow? This behaviour

was observed with smallish ([20,256,256]) images and default hyperparameters.

[ahatamiz]

Hi @Pyromonk

While the UNETR may have larger memory footprint comparing to CNNs, we have been able to handle most applications without any issues. Would you please provide more information regarding the input data ? I am assuming gray-scale images of size 256 x 256 with batch size of 20 ?

[Pyromonk]

Hi @Pyromonk

While the UNETR may have larger memory footprint comparing to CNNs, we have been able to handle most applications without any issues. Would you please provide more information regarding the input data ? I am assuming gray-scale images of size 256 x 256 with batch size of 20 ?

Hi, @ahatamiz!

My apologies for the somewhat late response.

But of course! It's reported as (8, 3, 20, 256, 256) by the DataLoader, so batch size 8, 3 channels (it is greyscale data though), image size 20x256x256 (these are 3D MRI scans).

I found the memory allocation problem very strange. The same input works perfectly fine with the baseline UNet model from MONAI, I see no reason for it to cause problems with UNETR.