Can't resolve OOM Error caused by FRNN

[mbendjilali]

Hello, I've been trying to run a training on dales using the command line

python src/train.py experiment=dales_11g ckpt_path=../dales_11g.ckpt

on a 2080Ti GPU with 11 Gb of RAM, but I always end up crashing because of a torch.cuda.OutOfMemoryError. I tried to tweak some of the parameters proposed in the README, but nothing does it. From my perspective, it looks like the training script always crash when calling frnn_grid_points. Here is an exemple traceback :

Error executing job with overrides: ['experiment=dales_11g', 'ckpt_path=../dales_11g.ckpt', 'datamodule.xy_tiling=6', 'datamodule.graph_k_max=2', 'datamodule.dataloader.batch_size=8']
Traceback (most recent call last):
  File "src/train.py", line 139, in main
    metric_dict, _ = train(cfg)
  File "/data/Moussa/spt/superpoint_transformer/src/utils/utils.py", line 48, in wrap
    raise ex
  File "/data/Moussa/spt/superpoint_transformer/src/utils/utils.py", line 45, in wrap
    metric_dict, object_dict = task_func(cfg=cfg)
  File "src/train.py", line 114, in train
    trainer.fit(model=model, datamodule=datamodule, ckpt_path=cfg.get("ckpt_path"))
  File "/home/moussa/miniconda3/envs/spt/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 532, in fit
    call._call_and_handle_interrupt(
  File "/home/moussa/miniconda3/envs/spt/lib/python3.8/site-packages/pytorch_lightning/trainer/call.py", line 43, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/home/moussa/miniconda3/envs/spt/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 571, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/moussa/miniconda3/envs/spt/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 932, in _run
    self._data_connector.prepare_data()
  File "/home/moussa/miniconda3/envs/spt/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/data_connector.py", line 94, in prepare_data
    call._call_lightning_datamodule_hook(trainer, "prepare_data")
  File "/home/moussa/miniconda3/envs/spt/lib/python3.8/site-packages/pytorch_lightning/trainer/call.py", line 166, in _call_lightning_datamodule_hook
    return fn(*args, **kwargs)
  File "/data/Moussa/spt/superpoint_transformer/src/datamodules/base.py", line 144, in prepare_data
    self.dataset_class(
  File "/data/Moussa/spt/superpoint_transformer/src/datasets/base.py", line 193, in __init__
    super().__init__(root, transform, pre_transform, pre_filter)
  File "/home/moussa/miniconda3/envs/spt/lib/python3.8/site-packages/torch_geometric/data/in_memory_dataset.py", line 57, in __init__
    super().__init__(root, transform, pre_transform, pre_filter, log)
  File "/home/moussa/miniconda3/envs/spt/lib/python3.8/site-packages/torch_geometric/data/dataset.py", line 97, in __init__
    self._process()
  File "/data/Moussa/spt/superpoint_transformer/src/datasets/base.py", line 493, in _process
    self.process()
  File "/data/Moussa/spt/superpoint_transformer/src/datasets/base.py", line 528, in process
    self._process_single_cloud(p)
  File "/data/Moussa/spt/superpoint_transformer/src/datasets/base.py", line 559, in _process_single_cloud
    nag = self.pre_transform(data)
  File "/home/moussa/miniconda3/envs/spt/lib/python3.8/site-packages/torch_geometric/transforms/compose.py", line 24, in __call__
    data = transform(data)
  File "/data/Moussa/spt/superpoint_transformer/src/transforms/transforms.py", line 23, in __call__
    return self._process(x)
  File "/data/Moussa/spt/superpoint_transformer/src/transforms/neighbors.py", line 46, in _process
    neighbors, distances = knn_1(
  File "/data/Moussa/spt/superpoint_transformer/src/utils/neighbors.py", line 36, in knn_1
    distances, neighbors, _, _ = frnn.frnn_grid_points(
  File "/data/Moussa/spt/superpoint_transformer/src/dependencies/FRNN/frnn/frnn.py", line 331, in frnn_grid_points
    idxs, dists, sorted_points2, pc2_grid_off, sorted_points2_idxs, grid_params_cuda = _frnn_grid_points.apply(
  File "/home/moussa/miniconda3/envs/spt/lib/python3.8/site-packages/torch/autograd/function.py", line 506, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/data/Moussa/spt/superpoint_transformer/src/dependencies/FRNN/frnn/frnn.py", line 174, in forward
    idxs, dists = _C.find_nbrs_cuda(sorted_points1, sorted_points2,
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 58.00 MiB (GPU 0; 10.75 GiB total capacity; 45.41 MiB already allocated; 44.50 MiB free; 62.00 MiB 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

[drprojects]

Hi @mbendjilali, did you make sure you had no other process running on your GPU ? In particular, are you perhaps using a desktop computer's GPU ? If so, you do not want any graphical user interface taking space on it, so make sure you are running on a GUI-free session. You do need to have all 11G 100% available for the python process.

Second, if that is not enough and FRNN is the source of memory error, you could try one of the following:

• increasing datamodule.xy_tiling=5 (will subdivide large tiles into 5x5 smaller tiles)
• reducing datamodule.knn=25 (number of neighbors for adjacency graph and geometric point features computation)
• reducing datamodule.knn_r=10 (maximum radius for KNN search)

Note that modifying these may produce different performance for the pretrained checkpoint you are using. So you may want to retrain your own model if you change parameters.

Hope that helps

[mbendjilali]

Thank you very much for your tips, that worked perfectly !

[drprojects]

Awesome ! Thanks for the feedback, I appreciate it 😊