手上有多块GPU，怎么提高3D GS的训练速度和渲染质量？

0. 这篇文章干了啥？

3D 高斯样条插值（3DGS）已成为3D新视角合成的一种流行技术，这主要归因于其相比之前方法（如NeRF）在训练和渲染速度上的优势。然而，大多数现有的3DGS流程受限于使用单个GPU进行训练，当应用于高分辨率或更大规模的场景时，会产生内存和计算瓶颈。例如，标准的Rubble数据集包含1657张图像，每张图像的分辨率为4K。单个A100 40GB GPU最多可以容纳1120万个高斯分布——远低于3DGS的质量饱和点。正如我们在第5.2.1节中所示，增加高斯分布的数量会持续改善重建质量。因此，为了扩大3DGS训练在参数数量和速度方面的规模，我们开发了Grendel分布式训练系统，该系统将3DGS训练分布在多个GPU上，并通过一个经验法则根据批量大小自动调整训练超参数。

当前处理大规模场景时绕过单设备3DGS训练管道限制的方法通常涉及压缩方案或选择性渲染。然而，这些方法仍然局限于单个GPU，并且在进一步扩展时受到类似的内存和计算瓶颈的根本限制。

随着大规模分布式机器学习成为行业标准，并成为许多最先进的AI模型（如大型语言模型）的基石，有理由期望跨多个GPU分配工作负载会带来类似的性能和规模提升，这使得3DGS能够处理更大、更高分辨率的场景并产生更高质量的输出。

然而，3DGS作为一个基于梯度但非基于神经网络的训练框架脱颖而出。它在训练迭代中具有动态和不平衡的工作负载模式，具有独特的计算管道。因此，现有的SysML框架和工具包并不适用于这个领域。针对神经网络的框架假设了一个由一系列常规张量操作（如GEMM）组成的一致且平衡的工作负载，这与3DGS训练和渲染所展现的动态、非均匀计算模式截然不同。因此，为3DGS设计一个可扩展的分布式系统具有独特的挑战性。

在本文中，我们提出了关于扩展高斯样条插值（Gaussian Splatting，简称GS）训练的几个关键观察点，这些观察点指导了我们为3DGS设计的分布式学习管道。例如，我们注意到在迭代过程中，3DGS训练管道的每个阶段都可以有效地并行化，但并行化的轴在各个阶段之间是不同的。在3DGS训练管道的混合并行化中，一些计算针对单个输出像素进行操作（允许像素级并行化），而其他计算则针对单个3D高斯分布进行操作（允许高斯级并行化）。这虽然实现了并行化，但需要在管道阶段之间进行通信。为了最小化通信，我们还观察到3DGS具有空间局部性，即只有少量高斯分布会影响每个图像块。最后，随着训练的进行，渲染输出像素的计算强度会发生变化，因此我们识别出动态且不平衡的工作负载。因此，任何静态的工作负载划分策略都将是次优的。

在本文中，我们描述了Grendel，这是一个基于这些观察点设计的分布式3DGS训练框架。Grendel使用高斯级分布——即，它将高斯分布在多个GPU之间进行分配——对于在训练迭代中表现出高斯级并行化的步骤，以及对于其他步骤使用像素级分布。它通过在像素级分布期间将连续的图像区域分配给GPU，并利用空间局部性来最小化在GPU之间传输的高斯分布数量，从而最小化了在高斯级和像素级分布之间切换时的通信开销。最后，Grendel采用了一个动态负载均衡器，该负载均衡器使用先前的训练迭代来分配像素级计算，以最小化工作负载的不平衡。

Grendel 还通过批量处理多张图像来扩大训练规模。这与传统的 3DGS 训练不同，后者仅使用批量大小为 1 的数据，这在我们分布式框架中会导致 GPU 利用率的降低。为了维持更大的批次中的数据效率和重建质量，需要重新调整优化器的超参数。为此，我们基于启发式独立梯度假设，为批量 3DGS 训练引入了一种自动超参数缩放规则。我们通过实验验证了所提方法的有效性——Grendel 支持使用大批量大小（我们测试了最多 32 个）进行分布式训练，同时与批量大小为 1 的情况相比，保持了重建质量和数据效率。

下面一起来阅读一下这项工作~

1. 论文信息

标题：On Scaling Up 3D Gaussian Splatting Training

作者：Hexu Zhao, Haoyang Weng, Daohan Lu, Ang Li, Jinyang Li, Aurojit Panda, Saining Xie

机构：纽约大学、太平洋西北国家实验室

原文链接：https://arxiv.org/abs/2406.18533

代码链接：https://github.com/nyu-systems/Grendel-GS

官方主页：https://daohanlu.github.io/scaling-up-3dgs/

2. 摘要

由于3D高斯样条插值（3DGS）在视觉质量和渲染速度方面的优越性，它在3D重建中越来越受到欢迎。然而，目前3DGS的训练只能在单个GPU上进行，由于内存限制，它处理高分辨率和大规模3D重建任务的能力受到了限制。我们介绍了Grendel，一个分布式系统，旨在划分3DGS参数并在多个GPU上并行计算。由于每个高斯分布仅影响渲染像素的一小部分动态子集，Grendel采用稀疏的全对全通信来将必要的高斯分布传输到像素分区，并执行动态负载均衡。与现有的一次使用一个相机视图图像进行训练的3DGS系统不同，Grendel支持多视图批量训练。我们探索了各种优化超参数缩放策略，并发现简单的sqrt(批量大小)缩放规则非常有效。使用大规模、高分辨率场景进行的评估表明，Grendel通过跨多个GPU扩展3DGS参数来提高渲染质量。在Rubble数据集上，我们通过在16个GPU上分布4040万个高斯分布，实现了27.28的测试PSNR，而单个GPU上使用1120万个高斯分布时，PSNR为26.28。Grendel是一个开源项目，可在https://daohanlu.github.io/scaling-up-3dgs/获取。

3. 效果展示

使用我们的分布式三维高斯渲染系统Grendel进行两次大规模、高分辨率的场景重建。这两个图像都使用16个GPU进行渲染。左图像和右图像分别使用4000万和2400万高斯表示。Grendel在这两个场景中都达到了最先进的质量（PSNR）。

4. 主要贡献

• 我们描述了 Grendel 的设计和实现，这是一个可扩展、内存高效、自适应的 3DGS 分布式训练系统。Grendel 允许批量 3DGS 训练扩展到最多 32 个 GPU 上运行。

• 我们探索了 3DGS 的大批量训练动态，确定了一个简单的 sqrt(batch_size) 学习率缩放策略，这使得在批量大小超过 1 的情况下能够进行有效的、无需超参数调整的训练。

• 我们展示了 Grendel 能够实现高分辨率大规模场景渲染：我们使用 16 个 GPU 和 MegaNERF中的大规模 Rubble 场景渲染 4K 图像。对于这个场景，Grendel 使用 4040 万个高斯分布实现了 27.28 的 PSNR，超过了当前最先进的方法。所需的内存超过了单个 GPU 的容量，这使得没有 Grendel 的技术很难以这种质量渲染这个场景。

5. 基本原理是啥？

图2(b)概述了Grendel的设计 。Grendel根据3DGS的混合并行性分配工作：它在高斯变换步骤中使用基于高斯分布的分配方式——即每个GPU对高斯分布的一个不相交子集进行操作，并在图像渲染和损失计算步骤中使用基于像素的分配方式——即每个GPU对像素的一个不相交子集进行操作。空间局部性特征允许Grendel在这些阶段之间转换时受益于稀疏的全对全通信。

基于高斯分布的分配 。Grendel将高斯分布（包括其参数和优化器状态）进行划分，并均匀地将它们分配给GPU。然后，每个GPU独立计算分配给它的3D高斯集合的高斯变换。我们发现，不同高斯分布所需的计算量并没有显著变化，因此，将高斯分布均匀分配给GPU可以让我们在适应最大数量高斯分布的同时，线性加速这一步骤的计算。

基于像素的分配 。我们在图像渲染和损失计算步骤中将连续的图像区域分配给GPU。分配连续区域允许我们利用空间局部性，并减少在这一步骤中GPU之间传输的高斯分布数量。在我们的实现中，我们通过将图像分成16×16像素的块，序列化这些块，然后使用自适应策略将连续的块序列分配给不同的GPU来划分批量中的每个图像。对于批处理，每个GPU可以从批量中的不同图像分配块，如图2(b)所示。

使用稀疏全对全通信传输高斯分布 。为了渲染图像像素，GPU需要访问与该像素相交的高斯分布，而这些高斯分布在训练过程中是依赖于视图的且会发生变化，因此无法预先确定。因此，Grendel在高斯变换之后包括了一个通信步骤。由于3DGS具有空间局部性，每个像素分区仅需要所有3D高斯分布的一个小子集。我们利用这一点来减少通信量：每个GPU首先确定用于渲染像素分区的相交高斯分布集合（如图4所示），然后使用稀疏全对全通信来检索与分区中任何像素相交的高斯分布。在反向传播过程中会进行反向的全对全通信。虽然Grendel的设计与用于分布式深度神经网络训练的FSDP有些相似，但存在一些重要的差异。首先，与FSDP中的权重分片不同，Grendel中的基于高斯分布的分配不仅是为了存储，还是为了计算（高斯变换）。其次，与FSDP使用密集的全收集通信传输权重分片不同，Grendel仅使用稀疏全对全通信传输相关的高斯分布。

6. 实验结果

为了检验独立梯度假设在实践中是否成立，我们分析了现实世界设置下梯度的平均每个参数方差。在图5中，我们绘制了从“Rubble”数据集上的预训练检查点开始的漫反射颜色参数的梯度的稀疏性和方差与批大小的关系。我们发现方差的倒数大致呈线性增加，然后过渡到平稳期。我们在三个检查点迭代中都发现了这种行为，分别代表训练初期、中期和后期。精度的初始线性增加表明，在本文工作中使用的批大小（高达32）下，梯度大致是不相关的，并支持独立梯度假设。然而，值得注意的是，尽管单个图像的梯度是稀疏的，但当一批图像很多时，梯度会重叠并变得不那么稀疏。同时，由于我们预计具有相似姿态的图像会提供相似的梯度，因此这些梯度也会变得更加相关。

为了实证检验所提出的学习率和动量缩放规则是否有效，我们使用批量大小为1的批次将“Rubble”场景训练至迭代15,000次。然后，我们重置Adam优化器的状态，并使用不同的批量大小继续训练。我们比较了当切换到更大批量大小时，不同学习率和动量缩放规则在维持相似训练轨迹方面的表现，如图6所示。由于3D GS的不同参数组具有极大的差异，我们专注于一个特定的组，即漫反射颜色，以使比较具有意义。图6a比较了三种不同的学习率缩放规则∈[constant, sqrt, linear]，其中只有我们提出的“sqrt”规则在不同训练批量大小下保持了较高的更新余弦相似度和相似的更新幅度。类似地，图6b显示了我们提出的指数动量缩放规则在更新余弦相似度上高于使动量系数保持不变的替代方案。

数据集。我们使用表1中所示的数据集和相应的分辨率设置来评估Grendel。其中，Rubble和MatrixCity Block_All代表了大多数现有3DGS系统无法达到的大规模数据集，而其他数据集在3DGS论文中常用。这些数据集在区域大小和分辨率上各不相同，以全面测试我们的系统。

我们使用Rubble场景来评估训练吞吐量。对于这次实验，我们使用了3500万个已训练至收敛的高斯分布。由于为这一场景渲染4K图像所需的时间较长，我们测量了训练另外10,000张图像的吞吐量，并在图7中报告了吞吐量（以每秒图像数计），同时改变GPU的数量（x轴）和批量大小（y轴）。我们观察到，由于内存需求，我们无法使用单个GPU渲染这一场景。此外，增加GPU数量和增加批量大小都能提高性能：在我们的实验中，性能从每秒5.55张图像（4个GPU，批量大小为1）提高到每秒38.03张图像（32个GPU，批量大小为64）。我们在附录C.2中提供了有关此实验的更多详细信息。

接下来，我们使用分辨率为980×545的Train场景来评估缩放过程中的图像质量（这张图像足够小，允许我们完全渲染图像）。我们在图8中的结果显示，使用相同总数的图像进行训练时，额外的GPU在保持图像质量的同时提高了吞吐量。值得注意的是，我们的16个GPU设置，在批量大小为32的情况下，在30K张图像上完成训练仅需2分钟42.97秒，这代表了我们所知的最新训练速度。

7. 总结 & 未来工作

我们已经描述了Grendel，一个分布式3D高斯训练系统。我们的评估表明，Grendel能够渲染现有方法无法处理的大规模场景的高分辨率高质量图像，并且展现出近乎完美的性能和内存可扩展性。我们计划将Grendel开源，以便社区可以更容易地对基于3DGS的大规模场景重建技术进行实验。

对更多实验结果和文章细节感兴趣的读者，可以阅读一下论文原文~

本文仅做学术分享，如有侵权，请联系删文。

计算机视觉工坊交流群

目前我们已经建立了3D视觉方向多个社群，包括 2D计算机视觉 、 大模型 、 工业3D视觉 、 SLAM 、 自动驾驶 、 三维重建 、 无人机 等方向，细分群包括：

2D计算机视觉： 图像分类/分割、目标/检测、医学影像、GAN、OCR、2D缺陷检测、遥感测绘、超分辨率、人脸检测、行为识别、模型量化剪枝、迁移学习、人体姿态估计等

大模型： NLP、CV、ASR、生成对抗大模型、强化学习大模型、对话大模型等

工业3D视觉： 相机标定、立体匹配、三维点云、结构光、机械臂抓取、缺陷检测、6D位姿估计、相位偏折术、Halcon、摄影测量、阵列相机、光度立体视觉等。

SLAM ： 视觉SLAM、激光SLAM、语义SLAM、滤波算法、多传感器融合、多传感器标定、动态SLAM、MOT SLAM、NeRF SLAM、机器人导航等。

自动驾驶： 深度估计、Transformer、毫米波|激光雷达|视觉摄像头传感器、多传感器标定、多传感器融合、自动驾驶综合群等、3D目标检测、路径规划、轨迹预测、3D点云分割、模型部署、车道线检测、BEV感知、Occupancy、目标跟踪、端到端自动驾驶等。

三维重建： 3DGS、NeRF、多视图几何、OpenMVS、MVSNet、colmap、纹理贴图等

无人机： 四旋翼建模、无人机飞控等

除了这些，还有 求职 、 硬件选型 、 视觉产品落地 、 最新论文 、 3D视觉最新产品 、 3D视觉行业新闻 等交流群

添加小助理: dddvision，备注： 研究方向+学校/公司+昵称 （如3D点云+清华+小草莓） , 拉你入群。

3D视觉学习知识星球

3D视觉从入门到精通知识星球 、国内成立最早、6000+成员交流学习。包括： 星球视频课程近20门（价值超6000） 、 项目对接 、 3D视觉学习路线总结 、 最新顶会论文&代码 、 3D视觉行业最新模组 、 3D视觉优质源码汇总 、 书籍推荐 、 编程基础&学习工具 、 实战项目 &作业 、 求职招聘&面经&面试题 等等。欢迎加入3D视觉从入门到精通知识星球，一起学习进步。