search

jeya-maria-jose / KiU-Net-pytorch

Official Pytorch Code of KiU-Net for Image/3D Segmentation - MICCAI 2020 (Oral), IEEE TMI
https://sites.google.com/view/kiunet/kiu-net
MIT License
357 stars 80 forks source link

High memory consumption #22

Closed mamba10-lab closed 3 years ago

mamba10-lab commented 3 years ago

Hi, thanks for the great work you've done. I am trying to use either kiunet or reskiunet models with images of size 512512, but it doesn't work with such size on VRAM=40gb with batch size =1. It works on smaller images (256256, batchsize = 3). It also takes too long to train it... I find it a little bit weird because the model doesn't have too mane parameters (only 300k), I trained much larger models faster (e.g. 3*10^6 parameters), so I am wondering if it is the expected behavior or maybe on new pytorch(i am using 1.7) it isn't efficient? Maybe you know some ways how I can improve the memory consumption and speed?

jeya-maria-jose commented 3 years ago

Hi,

Yeah. It does take a long time and memory to train for images for higher resolution. This happens as the feature map size in the bottle neck of Kite-Net branch is of high resolution and takes more compute power even though the number of parameters is less. Easy get-around to this problem is to use patch-wise training or remove some layers in the kite-net branch.

Any other fixes to this issue is welcome !

mamba10-lab commented 3 years ago

I will try that, thank you for the reply