Out-of-Memory Error in NeuS Optimization

[mshahbazi72]

Hi

I am trying to run the inference (including generation and reconstruction) on a RTX3090 GPU with 24GB memory.For some input images, I get Out-of-Memory error in the reconstruction part when setting --crop-size to 200 I. However, the code works fine on the same input image if I --crop-size to -1.

I am wondering what could be the reason, and if there is a way I can avoid OOO error in these cases.

Here is an example input image and the corresponding commands to reproduce the error:

python foreground_segment.py --input cup.png --output cup_fg.png

python generate.py --ckpt ckpt/syncdreamer-pretrain.ckpt \
                   --input cup_fg.png \
                   --output output/cup \
                   --sample_num 1 \
                   --cfg_scale 2.0 \
                   --elevation 30 \
                   --crop_size 200

python train_renderer.py -i output/cup/0.png \
                         -n cup-neus \
                         -b configs/neus.yaml \
                         -l output/renderer

[liuyuan-pal]

Hi, the crop size does not affect the amount of memory used in inference and the input size is always 256. There may be other factors that cause the OOM error.

[liuyuan-pal]

It seems that it guesses another handle on the backside sometimes. :)

[mshahbazi72]

I also can't see a connection between the option --crop-size and the memory issue. However, the memory error disappears if I decrease it from 200 to 150. I tried another image, and the same thing happened

[liuyuan-pal]

Maybe, the OOM does not occur on the generation but another process like foreground mask prediction.

[mshahbazi72]

Maybe, the OOM does not occur on the generation but another process like foreground mask prediction.

Yes, as I mentioned, it happens in the reconstruction part (training NeuS). The diffusion-based generation works fine. Could there be any correlation between the scale of the object in the image and the memory consumption during the training of NeuS?

[liuyuan-pal]

Oh, I know the reason. You may reduce the number of rays used in each iteration. Because only foreground rays will be fed into MLP. Larger foreground regions cause the OOM. https://github.com/liuyuan-pal/SyncDreamer/blob/eb41a0c73748cbb028ac9b007b11f8be70d09e48/configs/neus.yaml#L7 You need to reduce this number. This means we sample 3584+512 rays per batch with 3584 random rays and 512 foreground rays.

[mshahbazi72]

Thanks! This solved the problem.