共计 1697 个字符，预计需要花费 5 分钟才能阅读完成。
背景痛点
最近在做一个基于 3D 高斯泼溅 (3DGS) 的大规模点云重建项目，遇到一个棘手的问题：当点云规模达到百万级别时，训练过程中的显存占用会爆炸式增长。通过 NSight 工具分析发现，显存占用的峰值出现在反向传播阶段，主要原因是 Jacobian 矩阵的存储问题。
- 在标准 3DGS 实现中，每个高斯点的参数 (位置、颜色、协方差等) 都需要计算并存储完整的 Jacobian 矩阵
- 对于 100 万个点的场景，Jacobian 矩阵的显存占用轻松超过 10GB
- 反向传播时的中间激活值 (activations) 也需要大量显存，进一步加剧了问题
技术方案
经过调研和实验，我们最终采用了梯度检查点 (checkpointing) 和动态分辨率采样的混合方案。
梯度检查点 vs 激活值压缩
- 梯度检查点：只在关键层保存激活值，其他层在反向传播时重新计算
- 优点：显存节省显著(可达 50%)
-
缺点：增加约 30% 的计算时间
-
激活值压缩：使用有损压缩存储中间结果
- 优点：计算时间不变
- 缺点：可能影响模型精度
考虑到我们的场景对精度要求较高，最终选择了以梯度检查点为主的方案。
动态分辨率采样
为了进一步优化，我们引入了动态分辨率采样策略。核心思想是：
$$
\hat{S}t = S)
$$} \times \min(1, \frac{t}{T_{warmup}
其中：
– $S_{max}$ 是最大采样率
– $t$ 是当前训练步数
– $T_{warmup}$ 是预热步数
这个公式保证了训练初期使用低分辨率采样减轻负担，后期逐步提升到完整分辨率。
代码实现
CustomBackwardFunction
class CheckpointedGaussianBackward(torch.autograd.Function):
@staticmethod
def forward(ctx, inputs):
# 只保存必要的张量
ctx.save_for_backward(inputs)
return inputs
@staticmethod
def backward(ctx, grad_output):
inputs, = ctx.saved_tensors
# 重新计算中间结果
with torch.enable_grad():
outputs = gaussian_forward(inputs)
return grad_output * outputs多尺度 Dataloader
class MultiScaleLoader:
def __init__(self, base_loader):
self.loader = base_loader
self.scale = 0.1 # 初始采样率
def __iter__(self):
for data in self.loader:
# 应用动态采样
points = data['points']
N = int(len(points) * self.scale)
idx = torch.randperm(len(points))[:N]
yield {**data, 'points': points[idx]}混合精度配置
# 在 PyTorch Lightning 中配置
trainer = pl.Trainer(
precision='16-mixed',
accelerator='gpu',
devices=4
)性能验证
我们在不同硬件上测试了优化前后的性能：
| 硬件 | 原方案(iter/s) | 优化后(iter/s) | 显存节省 |
|---|---|---|---|
| RTX 3090 | 1.2 | 1.8 | 42% |
| A100 40G | 2.1 | 3.5 | 38% |
通过 Open3D 可视化对比，优化方案在重建质量上没有明显下降。
避坑指南
在实施过程中，我们总结了几点重要经验：
- 计算图重建陷阱：使用梯度检查点时，要确保重计算部分的随机种子一致
- 多 GPU 同步：动态采样率需要在所有 GPU 上保持同步
- 学习率调整：采样率变化时，建议同步调整学习率
结论与展望
通过梯度检查点和动态采样策略，我们成功将大规模 3DGS 训练的显存需求降低了 40%，同时保持了重建质量。不过，随着点云密度增加到千万级，还有更多优化方向值得探索：
- 更智能的梯度检查点策略
- 基于点云密度的自适应采样
- 分布式训练优化
希望这些实践经验对遇到类似问题的开发者有所帮助。完整的实现代码已开源在 GitHub 上，欢迎交流讨论。
正文完
