一种全新的数据蒸馏方法来加速NeRF

神经辐射场 (Neural Radiance Field, or NeRF) [Mildenhall et al., ECCV, 2020] 开启了用神经网络表征三维场景的新范式。NeRF 这两年在学术界和工业界都很火热, 但 NeRF 一个比较大的缺点是, 渲染速度慢。虽然 NeRF 用的神经网络 (11 层的 MLP) 本身很小, 但是渲染一个像素需要采集一条光线上的很多点(上百个), 这导致渲染一张图的计算量非常大, 如下图所示: 用 PyTorch 在单张 NVIDIA V100 显卡测试, 渲染 400x400 的图片就需要 6.7s 的时间, 这显然不利于 NeRF 在业界落地 (例如各种 AR/VR 设备, meta universe 等)。  

学术界已有不少研究工作来加速 NeRF。比较流行的一种方式是, 给定训练好的 NeRF, 采用更高效的数据结构进行存储, 如 Sparse Voxel Octree [Yu et al., ICCV, 2021]. 尽管加速很可观 (如 [Yu et al., ICCV, 2021] 实现了 3000x 的渲染加速), 但这种数据结构也破坏了 NeRF 作为场景表征存储小的优点。譬如, 原始 NeRF 网络仅仅 2.4MB 大小就可以存储一个场景, 而采用 Sparse Voxel Octree 则需要 1.93GB [Yu et al., ICCV, 2021], 这显然难以在端上应用。   因此, 如何加速 NeRF 渲染并维持其存储小的优点 (简言之: 小且快), 仍然是当前的研究热点, 也是本文的动因。  

Arxiv: https://arxiv.org/abs/2203.17261

Code: https://github.com/snap-research/R2L

Webpage: https://snap-research.github.io/R2L/

核心方法  我们所提出的核心方法从整体范式上来说非常简单: 通过数据蒸馏将神经辐射场 (NeRF) 转化为神经光场(Neural Light Field, or NeLF) -- 从 NeRF 到 NeLF, 所以我们把方法命名为 R2L。   NeLF 与 NeRF 一样, 都可以作为一个场景的表征. 不同的是:  

NeRF 的输入是场景中的一个点 (该点的坐标 + 该点所在视线的方向), 输出是该点的 RGB 和不透明度。NeRF 网络的输出是中间结果, 并不是图片上的 RGB 值. 要想得到一个像素的 RGB 值, 需要对该像素对应光线上的很多点进行积分 (即 Alpha Compositing)。

而 NeLF 的输入是一条光线, 输出直接是该光线对应图片上像素值, 不需要 Alpha Compositing 这一步。

对于 Novel View Synthesis 这个任务来说, NeLF 的优势很明显: 速度快! 要得到一个像素的 RGB 只需要跑一次网络, 而 NeRF 则需要跑上百次。  

但它的缺点也很明显, 主要有两个缺点:    (1)NeLF 网络要拟合的目标函数比 NeRF 更难。这一点可以这么理解: 在一张图片上相邻两个像素的 RGB 可能突变 (因为遮挡), 而相邻两个像素的光线方向其实差别很小, 这就意味着, 这个函数的输入稍微变化一点, 输出可能剧变, 这种函数的不连续性强, 复杂度高. 相比之下, NeRF 表达的函数是空间中的点, 空间中的点由于物理世界的连续性, 相邻位置上 RGB 剧变的可能性小, 所以函数相对简单。   (2)同样一堆图片, 用来训练 NeLF 的话, 样本量会大幅降低. 一张图片, 长宽为 H, W, 用来训练 NeLF 的话样本量就是 H*W, 而训练 NeRF 样本量是 H*W*K (K 是 NeRF 中的一条光线上的采样点个数, 在 NeRF 原文中 K=256). 所以, 从 NeRF 到 NeLF 训练样本量会变为原来的 1/K, 这是很大的缩减。   神经网络有效, 通常需要有大量的训练数据。从 NeRF 变为 NeLF, 一方面要拟合的目标函数变复杂了, 同时样本量却减小了, 无疑雪上加霜. 如何解决这些问题呢?   为了解决上述问题(1), 我们需要用一个更深的网络来表征更复杂的函数, 所以在我们的文章中提出了一个 88 层的深度残差 MLP (deep residual MLP), 网络结构如下:  

这样的深层网络在之前 NeRF 相关的工作没有出现过 (之前的 NeRF 相关工作大多继承了原始 NeRF 文章中的网络结构, 小修小补)。为了能让它训练起来, 我们引入了残差结构的设计。这一点跟 ResNet 的思想一样, 本身并没有更多的创新, 但把这一点引入到 NeRF/NeLF 中, 据我们所知, 本文是第一篇工作。残差结构的引入很有必要, 因为深度网络没有残差结构基本训练不起来, 这一点在文中的消融实验中也得到了证实。   另一个值得注意的创新点是关于如何表征一条光线。理论上说, 一条光线用一个方向向量就可以确定, 但如果真的只用方向向量去表征, 就会出现上面说的 “输入很接近, 输出却可能剧变” 的情况, 这就无疑会给 NeLF 网络的学习带来困难。为了使得 NeLF 网络要学习的函数更容易一些, 我们需要增强输入的差别.   具体来说, 我们采用一条光线上采样的多个点的坐标 (如下图所示), 将其串联(concat) 起来成一个向量, 以此作为该光线的表征, 作为我们 NeLF 网络的输入。  

这种表征非常简单直接, 同时也很有效。在文中, 我们也展示了它比之前的 NeLF 工作中用到的其他表征 (例如 Plucker 坐标 [Sitzmann et al, NeurIPS, 2021]) 要更为有效。   为了解决上述问题(2), 我们使用了一个预训练好的 NeRF 模型来产生大量伪数据 (pseudo data)。具体来说, 当 NeRF 对一个场景学习完之后, 给定任意一个角度 (ray direction), NeRF 都能返回这个角度下的图片, 我们就把这些图片收集起来, 形成了很多 (origin, direction, RGB) triplets。这些 triplets 就是训练我们模型的数据, loss 函数是 mean squared error (MSE), 如下所示:  

在我们的实验中, 我们收集了 10k 张图片, 是原始数据集 (大概 100 张图片) 的 100 倍, 这些数据确保了有充足的样本去训练 NeLF。文中的消融实验也表明, 大量伪数据对性能至关重要 (6.9dB PSNR 提升)!   值得一提的是, 如果仅仅是用伪数据训练, 我们的模型最优也只能复制 teacher NeRF, 无法超越它。为了能超越, 我们在原始图片上再微调 (Finetune) 一下模型。这个操作被证明有非常显著的效果, 使得我们的模型可以显著超越 teacher NeRF。  实验效果  总的来说, 我们的模型在 NeRF Synthetic 数据集 (图片尺寸 400x400) 上实现了将近 30x 的加速, 并把 PSNR 大幅提升了 1.4dB, 比同类其他方法更加高效。  

视觉效果图对比如下, 可以看到, 相比于 NeRF, 我们的模型 (Ours-2, 即在原始数据上微调后的模型)有肉眼可见的提升, 且计算量仅仅是 NeRF 的 1/26。  

更多结果请参考我们的文章。代码已经开源: https://github.com/snap-research/R2L, 欢迎尝试!  总结与未来工作  本文提出了一种全新的数据蒸馏方法来加速 NeRF: 我们使用训练好的 NeRF 模型产生伪数据, 来训练提出的深度残差 NeLF 网络。该 NeLF 网络可以达到超过 NeRF 的渲染质量, 且实现将近 30x 加速, 并维持了存储小的优点。   未来工作方向: (1) 从 NeRF 中可以得到深度信息, 目前我们还没提供从 NeLF 网络中得到深度信息的方法, 这是不错的探索方向。(2) 如何用更少, 更高质量的伪数据 (譬如进行数据筛选) 来加速 NeLF 的训练也非常值得探索。  

审核编辑 ：李倩