注意:翻译需要符合中文语序、流程、通顺
[新闻] 2023/04/04。如果你正在使用 nerfacc <= 0.3.5
并希望迁移到最新版本 (nerfacc >= 0.5.0
),请查看 CHANGELOG 了解如何迁移。
NerfAcc 是一个用于训练和推理的 PyTorch NeRF 加速工具箱。它专注于在辐射场的体渲染管道中进行高效采样,适用于大多数 NeRFs 的通用插件。 通过对现有代码库的最小修改,Nerfacc 能在训练各类最新的 NeRF 论文中提供显著的加速。 它是纯 Python 接口,具有灵活的 API!
安装
依赖:请先安装 Pytorch。
最简单的方法是从 PyPI 安装。这样,它将在 第一次运行时(JIT)构建 CUDA 代码。
pip install nerfacc
或者从源码安装。这样它将在安装过程中构建 CUDA 代码。
pip install git+https://github.com/nerfstudio-project/nerfacc.git
我们还提供了预构建的轮子,覆盖了 官方 Pytorch 支持的 Pytorch + CUDA 主要组合。
# 例如,torch 1.13.0 + cu117
pip install nerfacc -f https://nerfacc-bucket.s3.us-west-2.amazonaws.com/whl/torch-1.13.0_cu117.html
Windows 和 Linux | cu113 | cu115 | cu116 | cu117 | cu118 |
---|---|---|---|---|---|
torch 1.11.0 | ✅ | ✅ | |||
torch 1.12.0 | ✅ | ✅ | |||
torch 1.13.0 | ✅ | ✅ | |||
torch 2.0.0 | ✅ | ✅ |
对于之前版本的 nerfacc,请查看 这里 了解支持的预构建轮子。
使用方法
NerfAcc 的理念是使用计算成本低的估计器进行高效体积采样来发现表面。 因此,NerfAcc 可以与任何用户定义的辐射场一起工作。要将 NerfAcc 的渲染管道插入到你的代码中并享受加速效果,你只需要为你的辐射场定义两个函数。
sigma_fn
: 计算每个采样点的密度。它将被估计器(例如,nerfacc.OccGridEstimator
,nerfacc.PropNetEstimator
)用于发现表面。rgb_sigma_fn
: 计算每个采样点的颜色和密度。它将被nerfacc.rendering
用于进行可微体渲染。这个函数将接收梯度来更新你的辐射场。
一个简单的示例如下:
import torch
from torch import Tensor
import nerfacc
radiance_field = ... # 网络:一个 NeRF 模型
rays_o: Tensor = ... # 光线起点. (n_rays, 3)
rays_d: Tensor = ... # 光线标准化方向. (n_rays, 3)
optimizer = ... # 优化器
estimator = nerfacc.OccGridEstimator(...)
def sigma_fn(
t_starts: Tensor, t_ends:Tensor, ray_indices: Tensor
) -> Tensor:
""" 定义如何为估计器查询密度。"""
t_origins = rays_o[ray_indices] # (n_samples, 3)
t_dirs = rays_d[ray_indices] # (n_samples, 3)
positions = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0
sigmas = radiance_field.query_density(positions)
return sigmas # (n_samples,)
def rgb_sigma_fn(
t_starts: Tensor, t_ends: Tensor, ray_indices: Tensor
) -> Tuple[Tensor, Tensor]:
""" 从用户定义的辐射场中查询 rgb 和密度值。 """
t_origins = rays_o[ray_indices] # (n_samples, 3)
t_dirs = rays_d[ray_indices] # (n_samples, 3)
positions = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0
rgbs, sigmas = radiance_field(positions, condition=t_dirs)
return rgbs, sigmas # (n_samples, 3), (n_samples,)
# 高效光线行进:
# ray_indices: (n_samples,). t_starts: (n_samples,). t_ends: (n_samples,).
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,
)
# 可微体渲染:
# 颜色: (n_rays, 3). 不透明度: (n_rays, 1). 深度: (n_rays, 1).
color, opacity, depth, extras = nerfacc.rendering(
t_starts, t_ends, ray_indices, n_rays=rays_o.shape[0], rgb_sigma_fn=rgb_sigma_fn
)
# 优化: 网络和光线都会接收梯度
optimizer.zero_grad()
loss = F.mse_loss(color, color_gt)
loss.backward()
optimizer.step()
示例:
在运行这些示例脚本之前,请查看脚本需要哪个数据集,并首先下载数据集。你可以使用 --data_root
指定路径。
# 使用子模块克隆代码仓库.
git clone --recursive git://github.com/nerfstudio-project/nerfacc/
静态 NeRFs
完整的基准测试请见: https://www.nerfacc.com/en/stable/examples/static.html
在 NeRF-Synthetic 数据集上使用 Instant-NGP,在 4.5 分钟内获得更好的性能。
# 占据网格估计器
python examples/train_ngp_nerf_occ.py --scene lego --data_root data/nerf_synthetic
# 提案网估计器
python examples/train_ngp_nerf_prop.py --scene lego --data_root data/nerf_synthetic
在 Mip-NeRF 360 数据集上使用 Instant-NGP,在 5 分钟内获得更好的性能。
# 占据网格估计器
python examples/train_ngp_nerf_occ.py --scene garden --data_root data