nerfacc

repository·master·Indexed 23 days ago

https://github.com/nerfstudio-project/nerfacc

A PyTorch-based acceleration toolbox for NeRF (Neural Radiance Fields) training and inference. It provides a plug-and-play interface for existing radiance field implementations, focusing on efficient volumetric sampling and rendering using computationally cheap estimators like OccGridEstimator and PropNetEstimator. The library includes utility functions for raymarching, spatial intersections, and grid traversal, and is compatible with frameworks such as nerfstudio, sdfstudio, and instant-nsr-pl.

Tokens
6.1K
Snippets
6
Records
29
Agent score
76%

What's inside nerfacc

  1. Performance comparison of Camera Optimization NeRFs

    master
    The following table compares the performance of the BARF method against nerfacc using the occgrid (occupancy grid) implementation on the NeRF-Synthetic dataset. nerfacc demonstrates significantly faster training times and improved metrics (higher PSNR, lower LPIPS, $E_R$, and $E_T$) compared to the baseline BARF method.
  2. 3rd-Party integrations and use cases

    master

    The nerfacc library is designed to be compatible with several major NeRF and surface reconstruction frameworks, including:

    • nerfstudio: A collaboration-friendly studio for NeRFs.
    • modelscope: A collection of deep-learning algorithms.
    • sdfstudio: A unified framework for surface reconstruction.
    • instant-nsr-pl: A framework for training NeuS quickly.
  3. Radiance Field Implementation Details

    master
    The radiance field implementation in nerfacc follows the original NeRF paper's principles but optimizes the sampling strategy. Instead of the standard two-stage (coarse and fine) MLP sampling, it uses a single MLP with 1024 samples. This is made efficient by the library's fast rendering capabilities, which inherently skip samples that are far from the surface.
  4. Avoid CPU-GPU synchronization in PyTorch

    master

    CPU-GPU synchronization occurs when the CPU needs to know information that only the GPU currently holds (e.g., the number of True values in a boolean mask to determine the size of a new tensor). This forces the CPU to wait for the GPU, significantly increasing latency.

    Common operations that trigger synchronization:

    • torch.where
    • tensor.item()
    • print(tensor)
    • tensor.to(device)
    • torch.nonzero
    • Operations where the output shape depends on GPU-side data (e.g., boolean masking).

    Note for nerfacc users: Functions like nerfacc.traverse_grids and sampling with nerfacc.OccGridEstimator require synchronization because they must determine the output tensor size based on grid traversal results. This is an inherent requirement of these specific functions.

  5. How to integrate nerfacc into a radiance field pipeline

    master

    NerfAcc is designed to be plug-and-play for most NeRFs. To use the acceleration pipeline, you need to define two functions that interface with your radiance field:

    1. sigma_fn(t_starts, t_ends, ray_indices) -> Tensor: Computes density at each sample. This is used by estimators (like nerfacc.OccGridEstimator or nerfacc.PropNetEstimator) to discover surfaces.
    2. rgb_sigma_fn(t_starts, t_ends, ray_indices) -> Tuple[Tensor, Tensor]: Computes color and density at each sample. This is used by nerfacc.rendering for differentiable volumetric rendering. This function receives gradients to update your radiance field.

    Workflow

    1. Sampling: Use an estimator's .sampling() method to get ray_indices, t_starts, and t_ends.
    2. Rendering: Pass these values to nerfacc.rendering() to get color, opacity, and depth.
    3. Optimization: Perform standard PyTorch backpropagation. Both the network and the rays will receive gradients.
    import torch
    from torch import Tensor
    import nerfacc 
    
    # ... setup radiance_field, rays_o, rays_d, optimizer ...
    
    estimator = nerfacc.OccGridEstimator(...)
    
    def sigma_fn(t_starts: Tensor, t_ends: Tensor, ray_indices: Tensor) -> Tensor:
        """ Define how to query density for the estimator."""
        t_origins = rays_o[ray_indices]
        t_dirs = rays_d[ray_indices]
        positions = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0
        sigmas = radiance_field.query_density(positions) 
        return sigmas
    
    def rgb_sigma_fn(t_starts: Tensor, t_ends: Tensor, ray_indices: Tensor) -> Tuple[Tensor, Tensor]:
        """ Query rgb and density values from a user-defined radiance field. """
        t_origins = rays_o[ray_indices]
        t_dirs = rays_d[ray_indices]
        positions = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0
        rgbs, sigmas = radiance_field(positions, condition=t_dirs)  
        return rgbs, sigmas
    
    # 1. Efficient Raymarching
    ray_indices, t_starts, t_ends = estimator.sampling(
        rays_o, rays_d, sigma_fn=sigma_fn, near_plane=0.2, far_plane=1.0, early_stop_eps=1e-4, alpha_thre=1e-2
    )
    
    # 2. Differentiable Volumetric Rendering
    color, opacity, depth, extras = nerfacc.rendering(
        t_starts, t_ends, ray_indices, n_rays=rays_o.shape[0], rgb_sigma_fn=rgb_sigma_fn
    )
    
    # 3. Optimize
    optimizer.zero_grad()
    loss = F.mse_loss(color, color_gt)
    loss.backward()
    optimizer.step()
  6. Using nerfacc estimators with Dynamic NeRFs

    master

    When working with dynamic NeRFs (scenes that change over time), you can use nerfacc estimators to improve efficiency. The choice of estimator affects how temporal data is handled:

    • nerfacc.PropNetEstimator: Works naturally with dynamic NeRFs without additional modification.
    • nerfacc.OccGridEstimator: To use this with dynamic scenes, the recommended approach is to estimate the maximum opacity at each area over all timestamps. This effectively caches the union of occupancy across the entire temporal sequence. This allows you to share a single estimator across all timestamps (including those not in the training set), making rendering very efficient if motion is not extremely significant.
  7. Combine Instant-NGP with Proposal Networks

    master
    You can combine the Instant-NGP radiance field with a proposal network using nerfacc.PropNetEstimator. This approach (referred to as Ours (prop) in benchmarks) is functionally equivalent to the Nerfacto model used in the nerfstudio project. This configuration is useful for achieving high performance on datasets like Mip-NeRF 360.
  8. How NerfAcc works with user-defined radiance fields

    master

    NerfAcc accelerates the volumetric rendering pipeline by using a computationally cheap estimator to discover surfaces. To integrate NerfAcc with your existing radiance field, you must define two specific functions:

    1. sigma_fn: Computes density at each sample. This is used by an estimator (like nerfacc.OccGridEstimator or nerfacc.PropNetEstimator) to discover surfaces.
    2. rgb_sigma_fn: Computes both color and density at each sample. This is used by nerfacc.rendering for differentiable volumetric rendering. This function is part of the gradient flow used to update your radiance field.

    Once these are defined, you follow a two-step pipeline:

    1. Efficient Raymarching: Use an estimator's .sampling() method to get ray_indices, t_starts, and t_ends.
    2. Differentiable Volumetric Rendering: Pass those sampling results into nerfacc.rendering() to get colors, opacity, and depth.
  9. Maximize GPU utilization for performance

    master

    To achieve optimal performance, aim for 100% GPU utilization as reported by nvidia-smi. Low utilization typically stems from two causes:

    1. Insufficient Parallelism: The workload is too small to saturate the GPU kernels. Increasing the batch size can help.
    2. CPU-GPU Synchronization: The CPU is waiting for the GPU to finish a task, or the GPU is waiting for the CPU to provide information (like tensor shapes).

    Additionally, ensure that data loading and preprocessing do not bottleneck the GPU by using torch.utils.data.DataLoader to overlap data processing with GPU computation.

  10. Choose between OccGridEstimator and PropNetEstimator

    master

    When implementing efficient sampling in nerfacc, choose an estimator based on your scene type and performance requirements:

    EstimatorBest ForPros/Cons
    nerfacc.OccGridEstimatorScenes with mostly empty space (e.g., NeRF-Synthetic)Very efficient via voxel skipping, but may still place samples in occluded areas that contribute little.
    nerfacc.PropNetEstimatorUnbounded scenes (e.g., Mip-NeRF 360)More accurate transmittance estimation; concentrates samples on surfaces; works without a bounding box.

    Summary Decision Logic:

    • If your scene is bounded and mostly empty $\rightarrow$ Use OccGridEstimator.
    • If your scene is unbounded or requires high-precision surface sampling $\rightarrow$ Use PropNetEstimator.
  11. Implement a basic time-conditioned NeRF (T-NeRF)

    master

    For dynamic scene reconstruction, you can implement a time-conditioned NeRF (T-NeRF) model. The reference implementation follows the approach described in the D-NeRF paper, utilizing two MLPs:

    1. Radiance Field: An 8-layer MLP.
    2. Warping Field: A 4-layer MLP.

    Optimization Note: To account for relatively smooth object motion, the implementation reduces the maximum frequency of the positional encoding from 10 to 4 compared to standard NeRF implementations.

    Detailed implementation can be found in examples/radiance_fields/mlp.py or examples/train_mlp_dnerf.py in the repository.

  12. Install nerfacc

    master

    You can install nerfacc via PyPI, from source, or using pre-built wheels.

    Note: You must install PyTorch before installing nerfacc.

    PyPI (JIT Compilation)

    Installing from PyPI will build the CUDA code on the first run using Just-In-Time (JIT) compilation.

    pip install nerfacc

    From Source

    Installing from source builds the CUDA code during the installation process.

    pip install git+https://github.com/nerfstudio-project/nerfacc.git

    Pre-built Wheels

    To avoid JIT compilation, you can use pre-built wheels for specific PyTorch and CUDA combinations. For example, for torch 1.13.0 and cu117:

    pip install nerfacc -f https://nerfacc-bucket.s3.us-west-2.amazonaws.com/whl/torch-1.13.0_cu117.html
    pip install nerfacc