Asparagus15 / GaussianShader

code for GaussianShader: 3D Gaussian Splatting with Shading Functions for Reflective Surfaces
Other
288 stars 23 forks source link

How to lower memory consumption? #13

Closed cduguet closed 6 months ago

cduguet commented 6 months ago

Hi, after an hour of running the training script with a custom dataset of 224 images, I've got:

[ITER 7000] Evaluating test: L1 0.030716074630618095 PSNR 17.14025115966797 [27/02 16:35:09]

[ITER 7000] Evaluating train: L1 0.030992737039923668 PSNR 17.04730796813965 [27/02 16:35:14]

[ITER 7000] Saving Gaussians [27/02 16:35:14]
Training progress:  49%|██████████████████████████████████████████████████████████▍                                                             | 14620/30000 [1:25:05<1:31:25,  2.80it/s, Loss=0.05720Training progress:  49%|██████████████████████████████████████████████████████████▍                                                             | 14620/30000 [1:25:08<1:31:25,  2.80it/s, Loss=0.05952TrainiTraining progress:  61%|█████████████████████████████████████████████████████████████████████████▏                                              | 18310/30000 [1:47:45<1:13:37,  2.65it/s, Loss=0.0402723]Traceback (most recent call last):
  File "train.py", line 247, in <module>
    training(lp.extract(args), op.extract(args), pp.extract(args), args.test_iterations, args.save_iterations)
  File "train.py", line 83, in training
    render_pkg = render(viewpoint_cam, gaussians, pipe, background, debug=False)
  File "/data/projects/GaussianShader/gaussian_renderer/__init__.py", line 247, in render
    cov3D_precomp = cov3D_precomp)[0]
  File "/opt/conda/envs/gaussian_splatting/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/data/projects/GaussianShader/submodules/diff-gaussian-rasterization/diff_gaussian_rasterization/__init__.py", line 194, in forward
    raster_settings,
  File "/data/projects/GaussianShader/submodules/diff-gaussian-rasterization/diff_gaussian_rasterization/__init__.py", line 37, in rasterize_gaussians
    raster_settings,
  File "/data/projects/GaussianShader/submodules/diff-gaussian-rasterization/diff_gaussian_rasterization/__init__.py", line 78, in forward
    num_rendered, color, radii, geomBuffer, binningBuffer, imgBuffer = _C.rasterize_gaussians(*args)
RuntimeError: CUDA out of memory. Tried to allocate 1.60 GiB (GPU 0; 23.73 GiB total capacity; 11.84 GiB already allocated; 1.39 GiB free; 18.54 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
Training progress:  61%|█████████████████████████████████████████████████████████████████████████▏                                              | 18310/30000 [1:47:47<1:08:49,  2.83it/s, Loss=0.0402723]

I'm looking for a way to reduce the memory consumption, either in the rasterizer or in the input variables, so I don't have this problem. Any recommendations? Thank you!

cduguet commented 6 months ago

In my case it worked by actually following the warning recommendation, and using export 'PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512'