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:
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.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
- Sampling: Use an estimator's
.sampling() method to get ray_indices, t_starts, and t_ends. - Rendering: Pass these values to
nerfacc.rendering() to get color, opacity, and depth. - 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()