nerfacc - NeRF训练与推理的PyTorch加速工具箱

项目介绍：NerfAcc

NerfAcc 是一个基于 PyTorch 的工具箱，专门用于加速 NeRF（神经辐射场）的训练和推理过程。它的核心在于对辐射场的体渲染管线进行高效的采样，这种方法具备通用性，并能够与大多数 NeRF 无缝集成。用户只需对现有代码进行最小化修改，就可以大幅提高训练速度。NerfAcc 提供了灵活的 API，并且完全基于 Python 实现。

安装指南

使用 NerfAcc 前，需要先安装 PyTorch。安装 NerfAcc 的方法主要有两种：

从 PyPI 安装，首次运行时会即时编译 CUDA 代码。
```
pip install nerfacc
```
从源代码安装，此方式会在安装期间编译 CUDA 代码。
```
pip install git+https://github.com/nerfstudio-project/nerfacc.git
```

项目还提供了预构建的安装包，支持 Pytorch 与 CUDA 的多种组合。

使用方法

NerfAcc 旨在通过廉价的计算估计器进行高效的体积采样，进而发现表面。用户需要定义两个函数来利用 NerfAcc 的渲染能力：

sigma_fn：用于计算每个样本点的密度，这是发现表面的关键。
rgb_sigma_fn：用于计算每个样本点的颜色和密度，并将在渲染中接收梯度以更新辐射场。

一个简单的代码示例如下：

import torch
from torch import Tensor
import nerfacc 

radiance_field = ...
rays_o: Tensor = ...
rays_d: Tensor = ...
optimizer = ...

estimator = nerfacc.OccGridEstimator(...)

def sigma_fn(t_starts: Tensor, t_ends: Tensor, ray_indices: Tensor) -> Tensor:
    ...
    return sigmas

def rgb_sigma_fn(t_starts: Tensor, t_ends: Tensor, ray_indices: Tensor) -> Tuple[Tensor, Tensor]:
    ...
    return rgbs, sigmas

ray_indices, t_starts, t_ends = estimator.sampling(...)
color, opacity, depth, extras = nerfacc.rendering(...)
optimizer.zero_grad()
loss = F.mse_loss(color, color_gt)
loss.backward()
optimizer.step()

示例和应用

在使用示例脚本之前，请确保下载所需的数据集，并通过 --data_root 参数指定数据集路径。

静态 NeRFs

在 NeRF-Synthetic 数据集上改进的 Instant-NGP 实现，约需 4.5 分钟：

python examples/train_ngp_nerf_occ.py --scene lego --data_root data/nerf_synthetic

动态 NeRFs

在 D-NeRF 数据集上训练的 T-NeRF，约需 1 小时：

python examples/train_mlp_tnerf.py --scene lego --data_root data/dnerf

相机优化 NeRFs

在 NeRF-Synthetic 数据集上运行的 BARF：

cd benchmarks/barf/
bash script.sh nerfsyn-nerfacc-occgrid 0

常见安装问题

若遇到 ImportError，确保 Pytorch 和 CUDA 版本与 nerfacc 版本匹配。

论文引用

如果您使用了 NerfAcc，请引用以下论文：

@article{li2023nerfacc,
  title={NerfAcc: Efficient Sampling Accelerates NeRFs.},
  author={Li, Ruilong and Gao, Hang and Tancik, Matthew and Kanazawa, Angjoo},
  journal={arXiv preprint arXiv:2305.04966},
  year={2023}
}

NerfAcc 是一个强大而灵活的工具，为开发者和研究人员提供了方便高效的 NeRF 加速解决方案。通过合理使用，您将能显著提升工作效率，为各类神经辐射场的应用提供支持。