GPU OOM in val stage when without mask

[NNsauce]

hi, when I run "python launch.py --config configs/neus-dtu-wmask.yaml --gpu 1 --train", everything is ok, but when I run "python launch.py --config configs/neus-dtu.yaml --gpu 1 --train", it got CUDA out of memory . I am using the latest code where you've modify the chunk_batch function in models/utils.py as you said "move all output tensors to cpu before merging". I even set dynamic_ray_sampling=false or reduce max_train_num_rays to 2048, but the CUDA out of memory still happens.Could you please give me some advice, thx!!

[NNsauce]

here is complete error:

Epoch 0: : 0it [00:00, ?it/s]Update finite_difference_eps to 0.027204705103003882 Epoch 0: : 500it [00:26, 18.89it/s, loss=0.0754, train/inv_s=42.50, train/num_rays=512.0] Traceback (most recent call last): | 0/49 [00:00<?, ?it/s] File "launch.py", line 125, in main() File "launch.py", line 114, in main trainer.fit(system, datamodule=dm) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 608, in fit call._call_and_handle_interrupt( File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/trainer/call.py", line 36, in _call_and_handle_interrupt return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 88, in launch return function(*args, *kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 650, in _fit_impl self._run(model, ckpt_path=self.ckpt_path) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1112, in _run results = self._run_stage() File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1191, in _run_stage self._run_train() File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1214, in _run_train self.fit_loop.run() File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 199, in run self.advance(args, kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/loops/fit_loop.py", line 267, in advance self._outputs = self.epoch_loop.run(self._data_fetcher) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 200, in run self.on_advance_end() File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 250, in on_advance_end self._run_validation() File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 308, in _run_validation self.val_loop.run() File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 199, in run self.advance(*args, kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 152, in advance dl_outputs = self.epoch_loop.run(self._data_fetcher, dl_max_batches, kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 199, in run self.advance(args, kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 137, in advance output = self._evaluation_step(kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 234, in _evaluation_step output = self.trainer._call_strategy_hook(hook_name, kwargs.values()) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1494, in _call_strategy_hook output = fn(*args, *kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/strategies/ddp.py", line 359, in validation_step return self.model(args, kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(*input, kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 1008, in forward output = self._run_ddp_forward(*inputs, *kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 969, in _run_ddp_forward return module_to_run(inputs[0], kwargs[0]) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(*input, kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/pytorch_lightning/overrides/base.py", line 110, in forward return self._forward_module.validation_step(*inputs, *kwargs) File "/home/fzx/work/instant-nsr-pl/systems/neus.py", line 172, in validation_step out = self(batch) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(input, kwargs) File "/home/fzx/work/instant-nsr-pl/systems/neus.py", line 32, in forward return self.model(batch['rays']) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(*input, kwargs) File "/home/fzx/work/instant-nsr-pl/models/neus.py", line 293, in forward out = chunkbatch(self.forward, self.config.ray_chunk, True, rays) File "/home/fzx/work/instant-nsr-pl/models/utils.py", line 24, in chunk_batch out_chunk = func(*[arg[i:i+chunksize] if isinstance(arg, torch.Tensor) else arg for arg in args], **kwargs) File "/home/fzx/work/instant-nsr-pl/models/neus.py", line 230, in forward sdf, sdf_grad, feature, sdf_laplace = self.geometry(positions, with_grad=True, with_feature=True, with_laplace=True) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(*input, kwargs) File "/home/fzx/work/instant-nsr-pl/models/geometry.py", line 195, in forward points_d_sdf = self.network(self.encoding(points_d.view(-1, 3)))[...,0].view(points.shape[:-1], 6).float() File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(input, kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/torch/amp/autocast_mode.py", line 12, in decorate_autocast return func(*args, *kwargs) File "/home/fzx/work/instant-nsr-pl/models/network_utils.py", line 110, in forward x = self.layers(x.float()) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(input, kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/torch/nn/modules/container.py", line 139, in forward input = module(input) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(*input, **kwargs) File "/home/fzx/mambaforge/envs/sdf/lib/python3.8/site-packages/torch/nn/modules/activation.py", line 838, in forward return F.softplus(input, self.beta, self.threshold) RuntimeError: CUDA out of memory. Tried to allocate 1.19 GiB (GPU 0; 7.80 GiB total capacity; 3.63 GiB already allocated; 1.17 GiB free; 4.81 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 Epoch 0: : 500it [00:29, 16.70it/s, loss=0.0754, train/inv_s=42.50, train/num_rays=512.0]

[NNsauce]

I uh reduce chunk_size from 2048 to 1024, then it works. But why does it need so much more GPU memory when without mask?