ECCV 2022｜R2L: 用数据蒸馏减速NeRF

神经辐射场 (Neural Radiance Field, or NeRF) [Mildenhall et al., ECCV, 2020] 开启了用神经收集表征三维场景的新范式。NeRF 这两年在学术界和工业界都很火热, 但 NeRF 一个比较大的缺点是, 衬着速度慢。虽然 NeRF 用的神经收集 (11 层的 MLP) 本身很小, 但是衬着一个像素须要采集一条光芒上的很多点(上百个), 这导致衬着一张图的计算量非常大, 以下图所示: 用 PyTorch 在单张 NVIDIA V100 显卡测试, 衬着 400×400 的图片就须要 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.17261Code: https://github.com/snap-research/R2LWebpage: 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 数据集 (图片尺寸 400×400) 上实现了将近 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 的训练也非常值得探索。