NerfAcc: 加速NeRF训练与推理的PyTorch工具箱

NerfAcc: 让NeRF训练和推理更快捷高效

神经辐射场(Neural Radiance Fields, NeRF)技术近年来在3D场景重建和新视角合成领域取得了巨大突破。然而,NeRF模型的训练和推理过程往往十分耗时,这在很大程度上限制了其实际应用。为了解决这个问题,来自加州大学伯克利分校的研究人员开发了NerfAcc - 一个专门用于加速NeRF训练和推理的PyTorch工具箱。

NerfAcc的核心思想

NerfAcc的核心思想是通过高效的体积采样来加速NeRF的渲染过程。它引入了一种计算成本较低的估计器,用于发现3D场景中的表面。这种方法可以通用地应用于大多数NeRF变体,只需对现有代码进行最小的修改即可获得显著的速度提升。

NerfAcc的主要优势包括:

1. 通用性强: 可以应用于大多数NeRF变体
2. 易于集成: 只需对现有代码进行少量修改
3. 显著加速: 可以大幅提升训练和推理速度
4. 纯Python接口: 提供灵活易用的API

NerfAcc的工作原理

NerfAcc主要通过两个关键函数来实现高效渲染:

1. sigma_fn: 用于计算每个采样点的密度。这个函数被估计器(如OccGridEstimator或PropNetEstimator)用来发现场景中的表面。

2. rgb_sigma_fn: 用于计算每个采样点的颜色和密度。这个函数被nerfacc.rendering用来进行可微分的体积渲染。

通过这两个函数,NerfAcc可以高效地进行光线行进和体积渲染,从而大幅提升NeRF的训练和推理速度。

NerfAcc的安装和使用

NerfAcc的安装非常简单,可以通过pip直接安装:

pip install nerfacc

对于使用CUDA加速的用户,NerfAcc还提供了预编译的wheel包,覆盖了主流的PyTorch和CUDA版本组合。

使用NerfAcc只需要对现有的NeRF代码进行少量修改。以下是一个简单的示例:

import torch
import nerfacc

# 定义NeRF模型和光线
radiance_field = ...  # NeRF模型
rays_o = ...  # 光线原点
rays_d = ...  # 光线方向

# 定义密度函数和RGB-密度函数
def sigma_fn(t_starts, t_ends, ray_indices):
    # 计算采样点的密度
    ...
    return sigmas

def rgb_sigma_fn(t_starts, t_ends, ray_indices):
    # 计算采样点的颜色和密度
    ...
    return rgbs, sigmas

# 使用NerfAcc进行高效光线行进
estimator = nerfacc.OccGridEstimator(...)
ray_indices, t_starts, t_ends = estimator.sampling(
    rays_o, rays_d, sigma_fn=sigma_fn, 
    near_plane=0.2, far_plane=1.0
)

# 使用NerfAcc进行可微分体积渲染
color, opacity, depth, extras = nerfacc.rendering(
    t_starts, t_ends, ray_indices, 
    n_rays=rays_o.shape[0], rgb_sigma_fn=rgb_sigma_fn
)

通过这种方式,NerfAcc可以无缝集成到现有的NeRF项目中,为训练和推理过程带来显著的速度提升。

NerfAcc的实际应用效果

NerfAcc在多个benchmark上展现出了令人印象深刻的加速效果:

1. 在NeRF-Synthetic数据集上,使用Instant-NGP模型,NerfAcc可以在4.5分钟内达到与原始实现相当的性能。

2. 对于Mip-NeRF 360数据集,NerfAcc可以在5分钟内训练出高质量的模型。

3. 对于经典的MLP-based NeRF,NerfAcc可以将训练时间从原来的数天缩短到1小时左右。

4. 在动态NeRF场景中,NerfAcc同样表现出色。例如,它可以在1小时内训练出高质量的T-NeRF模型。

这些结果充分说明了NerfAcc在各种NeRF变体和数据集上的通用性和高效性。

NerfAcc的广泛应用

由于其出色的性能和易用性,NerfAcc已经被多个知名的3D视觉项目采用:

• nerfstudio: 一个用于协作开发NeRF的框架
• sdfstudio: 一个统一的表面重建框架
• threestudio: 一个用于3D内容创作的统一框架
• instant-nsr-pl: 可以在10分钟内训练NeuS模型

此外,NerfAcc还被应用于多个前沿的研究项目中,包括动态MLP地图表示、多视图人头重建等。

结语

NerfAcc为NeRF技术的实际应用扫清了一个重要障碍 - 训练和推理的时间成本。通过其高效的采样和渲染方法,NerfAcc使得高质量的3D场景重建和新视角合成可以在更短的时间内完成。这不仅为研究人员提供了更快的实验迭代周期,也为NeRF技术在实际应用中的部署铺平了道路。

随着NerfAcc的不断发展和完善,我们可以期待看到更多基于NeRF的创新应用出现在计算机视觉、增强现实、虚拟现实等领域。NerfAcc的出现,无疑为NeRF技术的未来发展注入了新的动力。