CUDA utils

[Kin-Zhang]

Chamfer CUDA utilitys comparsion, code path: https://github.com/KTH-RPL/OpenSceneFlow/blob/main/assets/tests/chamferdis_speed_test.py

Under 2080Ti:

(opensf) ➜  tests git:(main) ✗ python chamferdis_speed_test.py

Start status on GPU allocation: 0.000MB
torch.Size([88132, 3]) demo data:  tensor([-8.2266,  8.3516,  1.4922], device='cuda:0', grad_fn=<SelectBackward0>)
torch.Size([88101, 3]) demo data:  tensor([-7.9961,  8.1328,  0.4839], device='cuda:0', grad_fn=<SelectBackward0>)

Status after loading data: 2.017MB
------ START Pytorch3d Chamfer Distance Cal ------
Pytorch3d Chamfer Distance Cal time: 83.132 ms
loss:  tensor(0.1710, device='cuda:0', grad_fn=<AddBackward0>)

------ START CUDA Chamfer Distance Cal ------
loss:  tensor(0.1711, device='cuda:0', grad_fn=<AddBackward0>)
Chamfer Distance Cal time: 22.079 ms

Under 3090Ti:

0: 0.000MB
torch.Size([88132, 3]) demo data:  tensor([-8.2266,  8.3516,  1.4922], device='cuda:0', grad_fn=<SelectBackward0>)
torch.Size([88101, 3]) demo data:  tensor([-7.9961,  8.1328,  0.4839], device='cuda:0', grad_fn=<SelectBackward0>)
1: 2.017MB

------ START Faiss Chamfer Distance Cal ------
loss:  tensor(0.1710, device='cuda:0')
Faiss Chamfer Distance Cal time: 817.698 ms

------ START Pytorch3d Chamfer Distance Cal ------
Pytorch3d Chamfer Distance Cal time: 68.256 ms
loss:  tensor(0.1710, device='cuda:0', grad_fn=<AddBackward0>)

------ START mmcv 2D Chamfer Distance Cal ------
loss:  tensor(0.0591, device='cuda:0', grad_fn=<AddBackward0>)
mmcv Chamfer Distance Cal time: 651.510 ms

------ START CUDA Chamfer Distance Cal ------
Chamfer Distance Cal time: 14.308 ms
loss:  tensor(0.1710, device='cuda:0', grad_fn=<AddBackward0>)

[Kin-Zhang]

Some environment issue I met recently, mainly for custom CUDA utils.

I found out we cannot upgrade to CUDA-11.8, as there will be warning for torch_cpp and then the VRAM cost more than usual. Then raise up the illegal memory issue:

Traceback (most recent call last):
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 44, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 559, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 935, in _run
    results = self._run_stage()
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 978, in _run_stage
    self.fit_loop.run()
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/loops/fit_loop.py", line 201, in run
    self.advance()
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/loops/fit_loop.py", line 354, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 133, in run
    self.advance(data_fetcher)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 218, in advance
    batch_output = self.automatic_optimization.run(trainer.optimizers[0], kwargs)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 185, in run
    self._optimizer_step(kwargs.get("batch_idx", 0), closure)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 261, in _optimizer_step
    call._call_lightning_module_hook(
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 142, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/core/module.py", line 1265, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/core/optimizer.py", line 158, in step
    step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/strategies/strategy.py", line 224, in optimizer_step
    return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/plugins/precision/precision_plugin.py", line 114, in optimizer_step
    return optimizer.step(closure=closure, **kwargs)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/torch/optim/optimizer.py", line 280, in wrapper
    out = func(*args, **kwargs)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/torch/optim/optimizer.py", line 33, in _use_grad
    ret = func(self, *args, **kwargs)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/torch/optim/adam.py", line 121, in step
    loss = closure()
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/plugins/precision/precision_plugin.py", line 102, in _wrap_closure
    self._after_closure(model, optimizer)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/plugins/precision/precision_plugin.py", line 82, in _after_closure
    self._clip_gradients(
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/plugins/precision/precision_plugin.py", line 127, in _clip_gradients
    call._call_lightning_module_hook(
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 142, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/core/module.py", line 1192, in configure_gradient_clipping
    self.clip_gradients(
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/core/module.py", line 1164, in clip_gradients
    self.trainer.precision_plugin.clip_gradients(optimizer, gradient_clip_val, gradient_clip_algorithm)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/plugins/precision/precision_plugin.py", line 147, in clip_gradients
    self.clip_grad_by_norm(optimizer, clip_val)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/lightning/pytorch/plugins/precision/precision_plugin.py", line 157, in clip_grad_by_norm
    torch.nn.utils.clip_grad_norm_(parameters, clip_val)
  File "/home/kin/mambaforge/envs/opensf-test/lib/python3.10/site-packages/torch/nn/utils/clip_grad.py", line 76, in clip_grad_norm_
    torch._foreach_mul_(grads, clip_coef_clamped.to(device))  # type: ignore[call-overload]
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Then it cost me 2 hours try to fix it but fail. I noted down somethings here for future exploration on upgrading env.

1. It only affect custom CUDA so mainly for CUDA-chamfer? (not that sure)
2. It looks like torch things issue as the warning said:

/home/kin/mambaforge/envs/opensf-test-v2/lib/python3.8/site-packages/torch/utils/cpp_extension.py:398: UserWarning: There are no /home/kin/mambaforge/envs/opensf-test-v2/bin/x86_64-conda-linux-gnu-c++ version bounds defined for CUDA version 11.8

And then I check the file here: https://github.com/pytorch/pytorch/blob/main/torch/utils/cpp_extension.py#L58-L78

They didn't update the CUDA 11.8 and that might cost low-performance compile somehow I guess.

Previously I never found the problem till I trained the full and find out the VRAM cost too much than before.

While https://github.com/KTH-RPL/OpenSceneFlow/blob/3d53932d29ab0170181651f238b470d0637e2b8d/requirements.txt I created requirement.txt file previously for this one I didn't check whether it will have same problem etc.