CUDA OOM during TSDF

[lorenzobagnol]

Hi everyone,

I'm using fVDB Reality Capture 0.4.0 with fVDB-core v. 0.4.2 to reconstruct very complex meshes where I need to preserve a very high level of geometric detail.

My dataset consists of:

~600 high-quality images
Resolution: 6000×4000
Running on a cloud instance with an NVIDIA L4 GPU (24 GB VRAM)

To preserve reconstruction quality, I need to keep both:

a small voxel_size
a small truncation_margin

However, during the call to integrate_tsdf_with_features, I consistently hit a CUDA Out-Of-Memory (OOM) error.

Below is the relevant integration loop:

    enumerator = tqdm.tqdm(dataloader, unit="imgs", desc="Extracting TSDF") if show_progress else dataloader

    for i, tsdf_input in enumerate(enumerator):
        cam_to_world_matrix = camera_to_world_matrices[i].to(dtype=torch.float32, device=device)
        projection_matrix = projection_matrices[i].to(dtype=torch.float32, device=device)

        rgb_image, depth_image, weight_image = tsdf_input

        # Convert rgb image to feature dtype
        if feature_dtype == torch.uint8:
            rgb_image = (rgb_image * 255).to(feature_dtype)
        else:
            rgb_image = rgb_image.to(feature_dtype)

        # Prepare depth and weight tensors
        depth_image = depth_image.to(dtype)
        weight_image = weight_image.to(dtype)

        # Build the full feature tensor (B, H, W, 3) for integration
        rgb_for_integration = rgb_image

        accum_grid, tsdf, weights, features = accum_grid.integrate_tsdf_with_features(
            truncation_margin,
            projection_matrix.to(dtype),
            cam_to_world_matrix.to(dtype),
            tsdf,
            features,
            weights,
            depth_image.squeeze(0).to(device),
            rgb_for_integration.squeeze(0).to(device),
            weight_image.squeeze(0).to(device),
        )

        if show_progress:
            assert isinstance(enumerator, tqdm.tqdm)
            enumerator.set_postfix({"accumulated_voxels": accum_grid.num_voxels})

        # Prune out zero weight voxels to save memory
        new_grid = accum_grid.pruned_grid(weights > 0.0)
        tsdf = new_grid.inject_from(accum_grid, tsdf)
        features = new_grid.inject_from(accum_grid, features)
        weights = new_grid.inject_from(accum_grid, weights)
        accum_grid = new_grid

        del rgb_image, depth_image, weight_image
        torch.cuda.synchronize()
        torch.cuda.empty_cache()

Questions
Is this level of memory usage expected for this function with high-resolution images and dense voxel grids?
Are there recommended strategies to reduce VRAM usage while preserving reconstruction quality?
Would more aggressive voxel pruning help?
Right now I only prune voxels with weights > 0.0
Is there a way to not load everything on GPU?
I'm not using masks

At the moment it feels like memory usage grows excessively fast when using small voxel sizes and truncation margins.

Any suggestions would be greatly appreciated.

[matthewdcong]

I believe this should be fixed by #688 . Would you mind retrying with the latest ToT main and letting us know if it fixes the issue?

[harrism]

Verification of #688 vs #666 (synthetic inject_from loop, GPU)

I reproduced the leak mechanism directly. The reality-capture TSDF loop calls inject_from per image; #688 fixes a reference cycle through the autograd context in _InJectFn. Since that fix is 9 lines of pure Python in _dense.py, I A/B-tested it in place (no rebuild):

┌─────────────────────┬──────────────┬─────────┬───────────────────────────────┐
│    fvdb version     │  loop mode   │ grad_fn │     mem growth / 80 iters     │
├─────────────────────┼──────────────┼─────────┼───────────────────────────────┤
│ pre-#688            │ grad-tracked │ present │ +1831 MB (~23 MB/iter) — LEAK │
├─────────────────────┼──────────────┼─────────┼───────────────────────────────┤
│ pre-#688            │ no_grad      │ none    │ +0 MB                         │
├─────────────────────┼──────────────┼─────────┼───────────────────────────────┤
│ post-#688 (patched) │ grad-tracked │ present │ +0 MB — FIXED                 │
└─────────────────────┴──────────────┴─────────┴───────────────────────────────┘

#688 genuinely fixes the leak — the same grad-tracked loop that leaked ~23 MB/iteration leaks nothing after the patch. Over ~600 images that's ~14 GB of accumulated garbage, which squarely explains an OOM on a 24 GB L4. (It's a cyclic leak — gc.collect() reclaims it, but in a tight CUDA loop you OOM long before Python's GC runs.)

The caveat (crucial for #666)

The leak only triggers with autograd enabled. RC's stock tsdf_from_splats is decorated @torch.no_grad(), and under no_grad there's zero growth — so the stock library path never had this leak, and #688 wouldn't change its memory at all.

So #688 fixes #666 only if the OP's integration loop ran with grad tracking. Two tells point that way:

• The OP pasted a bare, extracted loop with no @torch.no_grad() — i.e. likely their own copy of the loop, running grad-tracked → leaks → Fix reference cycle in _InJectFn leading to a memory leak #688 fixes them.
• If instead they'd called the unmodified library function (no_grad), the OOM would be pure voxel-count (tiny voxel_size on a 6000×4000 × ~600 scene), and Fix reference cycle in _InJectFn leading to a memory leak #688 wouldn't help — that'd need a larger voxel_size/image_downsample_factor.

@lorenzobagnol was your TSDF loop wrapped in torch.no_grad()? If not (as their snippet suggests), #688 resolves it, and an immediate workaround on 0.4.x is to wrap the loop in with torch.no_grad().

I'm going to close this issue as fixed, but if the answer to the above question is YES, there might be a different memory leak so please re-open.

[lorenzobagnol]

Thanks a lot for looking into this. However, as you mentioned, the code snippet I posted was taken directly from the tsdf_from_splats_dlnr function in fvdb-Reality-Capture:

@torch.no_grad()
def tsdf_from_splats_dlnr(...)

So I don't think this particular leak is the cause of the OOM in my case. Also, I'm currently using fvdb-core v0.4.2 rather than the latest version. I cannot reopen this unfortunately.

Thanks again for your help. If you have any other ideas about what might be causing the memory issue, I'd really appreciate it.

[swahtz]

I've reopened it. Are you able to install fvdb-core 0.5.1 and see if this resolves memory issue @lorenzobagnol ?

[lorenzobagnol]

Thanks for reopening the issue.
I'm still using fvdb-core 0.4.2, as it's the version used by fVDB Reality Capture 0.4.0.
I've made quite a few modifications to my local Reality Capture codebase, so I won't be able to upgrade to 0.5.1 and test it anytime soon.

In the meantime, are there any other memory-related improvements between 0.4.2 and 0.5.1 (besides #688) that I could try to backport?

[swahtz]

If you can share voxel_size, truncation_margin, feature_dtype, and the accumulated_voxels count from the tqdm postfix at the point it OOMs, it might be possible to say whether you're hitting a transient peak or simply exceeding what 24 GB holds at that voxel size.

[lorenzobagnol]

I run another experiment using truncation_margin=0.2, grid_shell_thickness=4.0 and feature_dtype=uint8 on images of 3000x2000 pixels (so the same for the depth computed using the DLNR in fvdb-Reality-Capture).

When I call the integrate I obtained this CUDA OOM:
[00:48<00:10, 6.67imgs/s, accumulated_voxels=2.51e+8]CUDA error 2: out of memory (/__w/cpm_cache/nanovdb/0c4c067bd5f23f0c1b05ab224faa0ad89f73a4b9/nanovdb/nanovdb/cuda/DeviceBuffer.h:341)

At the end of the TSDF it has generated a very poor mesh with only 5,704 vertices and 8,235 faces. I checked the gaussian splat and it seems very good still.

Setup: NVIDIA L4 with 24GB.

[swahtz]

Hi @lorenzobagnol , it seems to me your voxel counts should be reasonable for that card. The peak memory usage for those numbers I'd have probably expected to be around 12-16GB. The quality issue of the mesh is likely to do with the process OOM'ing before the integration process has been able to reach completion.

The OOM issue you're probably running into is something we've been working a lot in NanoVDB recently to resolve which will land in #732. The issue is that the PyTorch caching allocator will likely have reserved memory that might be free but can't be used by cudaMalloc in NanoVDB, so the upstream changes that #732 integrates is to be able to pass PyTorch's allocator to NanoVDB for its memory allocations. You're welcome to check out that branch, build it and see if this resolves the issue, my guess is that this will help.