NerfStudio项目中torch.compile()对SplatFacto模型训练的影响分析

问题背景

在计算机视觉和3D重建领域，NerfStudio项目中的SplatFacto模型是一个重要的3D高斯泼溅(3D Gaussian Splatting)实现。近期有开发者发现，在该模型的get_viewmat()函数上使用PyTorch 2.0的@torch.compile()装饰器会导致训练结果陷入局部最优，而移除该装饰器后模型训练则表现正常。

技术细节分析

@torch.compile()是PyTorch 2.0引入的重要特性，旨在通过图编译优化模型执行效率。然而，在某些特定场景下，这种优化可能会带来意外的副作用：

1. 数值精度问题：图编译优化可能改变计算图的执行顺序，导致浮点运算的微小差异累积
2. 梯度计算变化：编译后的计算图可能以不同方式处理梯度传播
3. 硬件兼容性：不同CUDA版本和显卡驱动对编译优化的支持程度不一

实验观察

对比实验显示：

• 使用@torch.compile()时：

  • 训练PSNR值停滞在较低水平
  • 模型陷入局部最优，无法继续优化
  • 最终重建质量较差

• 不使用@torch.compile()时：

  • PSNR曲线呈现正常上升趋势
  • 模型能够收敛到更好的解
  • 重建质量符合预期

解决方案与建议

基于当前发现，我们建议：

1. 版本升级：确保使用PyTorch 2.1.2或更高版本，早期版本(如2.0.1)已知存在编译优化问题
2. 选择性编译：对于关键函数如get_viewmat()，暂时避免使用编译优化
3. 性能权衡：在模型稳定性和执行速度之间做出合理取舍
4. 环境一致性：保持CUDA驱动、PyTorch版本和硬件配置的兼容性

深入技术探讨

get_viewmat()函数负责计算视图变换矩阵，这一操作在3D重建中至关重要。当使用@torch.compile()时，可能导致：

• 视图矩阵计算的微小误差被放大
• 反向传播路径被优化改变
• 自动微分行为出现偏差

这些因素综合作用，最终影响了整个模型的优化轨迹。

结论

在计算机视觉和深度学习项目中，性能优化工具的使用需要谨慎评估。NerfStudio项目中SplatFacto模型的这一案例表明，即使是官方提供的优化手段，也可能在特定场景下产生负面影响。开发者应当在保证模型正确性的前提下，逐步引入性能优化措施，并通过严格的验证确保优化不会损害模型质量。