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

CUDA out of memory #1

Closed XiaoJiNu closed 4 years ago

XiaoJiNu commented 4 years ago

Hi, when I trained the model with custom data, it always told that CUDA out of memory. I tried resize image to 608x608 and use bath size 1 but can't solve the problem. What's wrong?

jeya-maria-jose commented 4 years ago

@XiaoJiNu So, this code works perfectly with image size as 128x128 using any GeForce version GPU with batch size 1. You can do either of the following: 1) resize the image to 128x128 or somewhere closer to that resolution. 2) Reduce the number of layers and filters in the network by modifying the code such that you can accommodate 608*608. 3) Do patchwise training instead. Just use the crop argument to feed in the size of the patch that you would like to train.