search

askerlee / segtran

Medical Image Segmentation using Squeeze-and-Expansion Transformers
216 stars 49 forks source link

CUDA out of memory. #7

Closed wshi8 closed 3 years ago

wshi8 commented 3 years ago

Hello,

I run this command with batch size as 2 on the sample data provided, and it got a "CUDA out of memory" error. Your advice is appreciated.

~/segtran/code$ python3 train3d.py --task brats --split all --bs 1 --maxiter 10000 --randscale 0.1 --net segtran --attractors 1024 --translayers 1

Traceback (most recent call last): File "train3d.py", line 683, in outputs = net(volume_batch) File "/home/ubuntu/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl return forward_call(*input, kwargs) File "/home/ubuntu/segtran/code/networks/segtran3d.py", line 541, in forward vfeat_fused_fpn = self.out_fpn_forward(batch_base_feats, vfeat_fused) File "/home/ubuntu/segtran/code/networks/segtran3d.py", line 428, in out_fpn_forward align_corners=False) File "/home/ubuntu/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/torch/nn/functional.py", line 3712, in interpolate return torch._C._nn.upsample_trilinear3d(input, output_size, align_corners, scale_factors) RuntimeError: CUDA out of memory. Tried to allocate 1.15 GiB (GPU 0; 15.78 GiB total capacity; 11.45 GiB already allocated; 900.00 MiB free; 13.54 GiB reserved in total by PyTorch)**

wshi8 commented 3 years ago

After I did changes based on the links below, I can run batch 2 now:

https://spell.ml/blog/mixed-precision-training-with-pytorch-Xuk7YBEAACAASJam https://pytorch.org/docs/stable/amp.html

askerlee commented 3 years ago

Thanks for providing your workaround. On my GPU this command line consumes 17.5GB of GPU RAM. My GPU is Titan RTX 24GB so it's not a problem...