search

junyuchen245 / TransMorph_Transformer_for_Medical_Image_Registration

TransMorph: Transformer for Unsupervised Medical Image Registration (PyTorch)
MIT License
437 stars 73 forks source link

Memory problems during model training and testing #38

Closed Yongxin1120 closed 1 year ago

Yongxin1120 commented 2 years ago

Hi Junyu

Thanks for your help. Recently we encountered some memory problems during model training and testing. As shown in the picture, on our workstation, we are using an NVIDIA RTX A2000 graphics card with 12GB of video memory, but it still prompts that there is not enough memory. I would like to ask what configuration of video card is required to run this model, and whether there is a way to reduce the amount of video memory needed to run this model.

Best regards

Yongxin

cece5c87b0e6a8c987423515c414d14
junyuchen245 commented 1 year ago

Hi @Yongxin1120 ,

Enabling the checkpointing flag can help save memory. Additionally, we will be providing a script that works on 16FP in the future. With the image size of the IXI dataset, the memory usage should be well under 10GB. https://github.com/junyuchen245/TransMorph_Transformer_for_Medical_Image_Registration/blob/9c34ac92ec4626c88714fd2336aaff59c163e626/IXI/TransMorph/models/configs_TransMorph.py#L51

Thanks, Junyu