史上最大重建25km²！NeRF-XL：真正有效利用多卡联合训练！

原标题：NeRF-XL: Scaling NeRFs with Multiple GPUs

论文链接：https://research.nvidia.com/labs/toronto-ai/nerfxl/assets/nerfxl.pdf

项目链接：https://research.nvidia.com/labs/toronto-ai/nerfxl/

作者单位：NVIDIA 加州大学伯克利分校

论文思路：

本文提出了NeRF-XL，这是一种原理性的方法，用于在多个图形处理器（GPUs）之间分配神经射线场（NeRFs），从而使得具有任意大容量的NeRF的训练和渲染成为可能。本文首先回顾了现有的几个GPU方法，这些方法将大型场景分解成多个独立训练的NeRFs [9, 15, 17]，并确定了这些方法的几个基本问题，这些问题在使用额外的计算资源（GPUs）进行训练时阻碍了重建质量的提高。NeRF-XL解决了这些问题，并允许通过简单地使用更多的硬件来训练和渲染具有任意数量参数的NeRFs。本文方法的核心是一个新颖的分布式训练和渲染公式，这在数学上等同于经典的单GPU案例，并最小化了GPU之间的通信。通过解锁具有任意大参数数量的NeRFs，本文的方法是第一个揭示NeRFs GPU扩展规律（scaling laws）的方法，显示出随着参数数量的增加而提高的重建质量，以及随着更多的GPU的增加而提高的速度。本文在多种数据集上展示了NeRF-XL的有效性，包括包含约258K张图像、覆盖了25平方公里的城市区域的MatrixCity [5]。

论文设计：

近期在新视角合成的进步极大地提高了我们捕获神经辐射场（NeRFs）的能力，使得这一过程变得更加易于接近。这些进步使得我们能够重建更大的场景和场景内更精细的细节。无论是通过增加空间规模（例如，捕获数公里长的城市景观）还是提高细节水平（例如，扫描田野中的草叶），扩大捕获场景的范围都涉及到将更多的信息量纳入NeRF中，以实现准确的重建。因此，对于信息含量高的场景，重建所需的可训练参数数量可能会超过单个GPU的内存容量。

本文提出了NeRF-XL，这是一个原理性的算法，用于在多个GPU之间高效分配神经渐射场景（NeRFs）。本文的方法通过简单增加硬件资源，使得捕获高信息含量的场景（包括大规模和高细节特征的场景）成为可能。NeRF-XL的核心是在一组不相交的空间区域之间分配NeRF参数，并跨GPU联合训练它们。不同于传统的分布式训练流程在后向传播中同步梯度，本文的方法只需要在前向传播中同步信息。此外，通过仔细渲染方程程和分布式设置中相关的损失项，本文大幅减少了GPU之间所需的数据传输。这种新颖的重写提高了训练和渲染的效率。本文方法的灵活性和可扩展性使得本文能够够够多个GPU高效地优化和使用多个GPU高效性能优化。

本文的工作与最近采用了GPU算法来建模大规模场景的方法形成了对比，这些方法通过训练一组独立立体的NeRFs来实现[9, 15, 17]。虽然这些方法不需要GPU之间的通信，但每个NeRF都需要建模整个空间，包括背景区域。这导致随着GPU数量的增加，模型容量中的冗余度增加。此外，这些方法在渲染时需要混合NeRFs，这会降低视觉质量并在重叠区域引入伪影。因此，与NeRF-XL不同的是，这些方法在训练中使用更多的模型参数（相当于更多的GPU时），未能实现视觉质量的提升。

本文通过一系列多样化的捕获案例来展示本文方法的有效性，包括街道扫描、无人机飞越和以物体为中心的视频。这些案例的范围从小场景（10平方米）到整个城市（25平方公里）。本文的实验表明，随着本文将更多的计算资源分配给优化过程，NeRF-XL开始能够实现改善的视觉质量（通过PSNR测量）和渲染速度。因此，NeRF-XL使得在任何空间规模和细节的场景上训练工具有任意容量的NeRF成为可能。

图 1：本文基于原理的多GPU分布式训练算法能够将NeRFs扩展到任意大的规模。

图 2：独立训练与多GPU联合训练。独立地训练多个NeRFs[9,15,18]要求每个NeRF既要建模焦点区域也要建模其周围环境，这导致了模型容量的冗余。相比之下，本文的联合训练方法使用不重叠的NeRFs，因此没有任何冗余。

图 3：独立训练需要在新视角合成时进行混合。无论是在2D[9, 15]还是3D[18]中进行混合，都会在渲染中引入模糊。

图 4：独立训练导致不同的相机优化。在NeRF中，相机优化可以通过变换不准确的相机本身或所有其他相机以及底层3D场景来实现。因此，伴随相机优化独立训练多个NeRF可能导致相机校正和场景几何的不一致性，这给混合渲染带来了更多困难。

图 5：3D混合可能造成的视觉伪影。左图展示了使用2个GPU训练的MegaNeRF结果。在0%重叠时，由于独立训练，边界出现了伪影；而在15%重叠时，由于3D混合，出现了严重的伪影。右图阐释了这种伪影的成因：虽然每个独立训练的NeRF渲染出正确的颜色，但混合后的NeRF并不保证正确的颜色渲染。

图 6：本文的训练流程。本文的方法联合训练所有GPU上的多个NeRFs，每个NeRF覆盖一个不相交的空间区域。GPU之间的通信仅发生在前向传播中，而不发生在后向传播中（如灰色箭头所示）。(a) 本文可以通过评估每个NeRF以获得样本颜色和密度，然后将这些值广播到所有其他GPU以进行全局体渲染（见第4.2节）。(b) 通过重写体渲染方程，本文可以将数据传输量大幅减少到每条光线一个值，从而提高效率（见第4.3节）。

实验结果：

图 7：定性比较。与先前的工作相比，本文的方法有效地利用多GPU配置，在所有类型的数据上提高了性能。

图 8：定量比较。基于独立训练的先前工作未能随着额外GPU的增加而实现性能提升，而本文的方法随着训练资源的增加，享受到了渲染质量和速度的提升。

图 9：本文方法的可扩展性。更多的GPU允许有更多的可学习参数，这导致了更大的模型容量和更好的质量。

图 10：大规模捕获上的更多渲染结果。本文在更大的捕获数据集上使用更多的GPU测试了本文方法的鲁棒性。请参阅本文的网页，以获取这些数据的视频导览。

图 11：在University4数据集上与PyTorch DDP的比较。PyTorch 分布式数据并行（Distributed Data Parallel，DDP）旨在通过跨GPU分布光线来加快渲染速度。相比之下，本文的方法是跨GPU分布参数，突破了集群中单个GPU的内存限制，并且能够扩大模型容量以获得更好的质量。

图 12：University4上的同步成本。本文基于分区的体渲染（见第4.3节）允许 tile-based 通信，这比原始的基于样本的通信（见第4.2节）成本要低得多，因此能够实现更快的渲染。

总结：

总结来说，本文重新审视了将大规模场景分解为独立训练的NeRFs（神经辐射场）的现有方法，并发现了阻碍额外计算资源（GPUs）有效利用的重大问题，这与利用多GPU设置来提升大规模NeRF性能的核心目标相矛盾。因此，本文引入了NeRF-XL，这是一种原理性的算法，能够有效地利用多GPU设置，并通过联合训练多个非重叠的NeRFs来在任何规模上增强NeRF性能。重要的是，本文的方法不依赖于任何启发式规则，并且在多GPU设置中遵循NeRF的扩展规律(scaling laws)，适用于各种类型的数据。

引用：

@misc{li2024nerfxl,title={NeRF-XL: Scaling NeRFs with Multiple GPUs}, author={Ruilong Li and Sanja Fidler and Angjoo Kanazawa and Francis Williams},year={2024},eprint={2404.16221},archivePrefix={arXiv},primaryClass={cs.CV}}