Jittor开源: JNeRF带你5秒训好NeRF

system · 2022 年6 月 1 日 13:35

计图(Jittor)框架的NeRF模型库JNeRF正式开源了！

通过JNeRF可以5秒训练好NeRF模型(见图1)。

Jittor成为首个支持Instant-NGP的深度学习框架。

图1 5秒训练NeRF效果展示

Part1 研究背景

NeRF在2020[1]被提出，它惊艳的视角生成效果获得了学术界的广泛关注。不同于传统的显式表达，NeRF通过将场景表示为隐式的神经辐射场，渲染时通过神经网络查询位置上的场景信息生成新视角图像。

1654089420(1)

图2 用Jittor框架训练的NeRF的例子

NeRF的出现很可能带来一场革命，但它仍存在一些问题。不同方法针对NeRF各部分(如：采样、位置编码、网络结构)的优化都会对NeRF的性能带来提升，这一定程度上影响了比较的公平性。并且随着NeRF效率的提升，深度学习框架逐渐成为NeRF效率提升的瓶颈，许多工作修改、添加cuda代码以提升效率，这进一步带来了混乱。

因此NeRF需要一个统一、高效、并与框架高度适配的模型库，建立统一流程，提升科研效率。

Part2 JNeRF架构

JNeRF模型库分析了现有的NeRF方法，将NeRF的主要训练流程总结为图3所示的7个模块，不同模块之间相互解耦，使其能够方便地调用和替换。由于时间有限，目前JNeRF只支持了其中很少的工作，后续JNeRF会不断支持更多有代表性的NeRF工作，也十分欢迎大家在github为JNeRF贡献代码。

图3 JNeRF模型库架构

Part3 JNeRF支持Instant-NGP

Nvidia的工作Instant-NGP[2]在今年年初发布后，受到了学术界和工业界的广泛关注，github开源仓库至今已获得7k+个星。

该工作通过Hash编码和定制化的优化，使得NeRF能在5秒训练出一个效果优质的结果。经过我们的分析，Instant-NGP能5秒训练出NeRF并不仅仅是因为哈希编码方法，更因为Nvidia针对硬件做的极致优化。

图4 Instant-NGP 算法流程

Instant-NGP 能5秒训练NeRF的效果给学术界和产业界都带来了很多可能，但Instant-NGP的实现仍存在一些问题：其源码完全基于cuda编写，对习惯python的用户而言使用门槛较高。Instant-NGP没有成熟的深度学习框架支持，无法适配各种常用NeRF变种模型；Nvidia为了追求极致的效率优化，其源码不同函数间耦合严重，难以修改。

基于JNeRF实现的Instant-NGP具有以下优势：

准确复现了论文的速度和点数(见表1)，其他深度学习框架复现的Instant-NGP在速度和点数上都距离原论文有一定差距，Jittor是全球唯一成功复现了Instant-NGP的框架。

表1 与Instant NGP原文的对比

JNeRF速度十分高效，其训练速度可以达到大约133 iter/s。我们注意到，随着NeRF训练速度的提升，框架的运行速度成为了限制NeRF速度进一步提升的瓶颈，近期有的工作（如Plenoxel）通过大量修改pytorch源码才实现快速训练，而Instant-NGP则更是直接使用cuda来达到速度的极致追求。得益于Jittor的速度优势，JNeRF可以直接轻松地使用Jittor框架完成Instant-NGP的复现。
Jittor是一个即时编译的框架，通过即时编译的特性，能够实现很多原版Instant-NGP都无法实现的操作，例如在训练时动态修改cuda代码。
JNeRF进行了较好的模块化，使得科研人员能够轻松地调用、替换其中的模块。
Jittor基于元算子融合的优势，有着丰富的硬件生态支持，因此基于Jittor的JNeRF库能够轻松地在各种国内外硬件以及移动端上运行。

Part4 JNeRF代码演示

JNeRF通过配置文件可以轻松修改模型设置，下面的代码展示了用户可以通过简单地修改配置文件实现模块的替换。

sampler = dict(
    type='DensityGirdSampler',
    update_den_freq=16,
)
encoder = dict(
    pos_encoder = dict(
        type='HashEncoder',
    ),
    dir_encoder = dict(
        type='SHEncoder',
    ),
)
model = dict(
    type='NGPNetworks',
    use_fully=True,
)
loss = dict(
    type='HuberLoss',
    delta=0.1,
)

目前，JNeRF仓库已在github和中国计算机学会的gitlink上开源。JNeRF仍有待不断完善，支持的模型仍十分有限，欢迎大家通过该仓库进行NeRF的学习和研究，并和我们一起完善该仓库，共同推动NeRF的发展。

Github链接：https://github.com/Jittor/JNeRF

Gitlink链接：https://www.gitlink.org.cn/jittor/jnerf

如果大家在使用过程中发现任何问题，或者希望一起推动该仓库的发展，请在github提交issue或者pr，也可以加入计图开发者交流群与我们讨论。

参考文献

Mildenhall B, Srinivasan P P, Tancik M, et al. Nerf: Representing scenes as neural radiance fields for view synthesis[C]//European conference on computer vision. Springer, Cham, 2020: 405-421.
MüllerT, Evans A, Schied C, et al., Instant Neural Graphics Primitives with a Multiresolution Hash Encoding, arXiv: 2201.05989, 2022.