Py学习  »  机器学习算法

TRO 开源新作:让 SLAM/SfM 核心后端 BA 跑进 PyTorch,几何优化与深度学习终于接上了

智驾机器人技术前线 • 3 天前 • 26 次点击  

点击下方卡片关注我们,点亮星标⭐,优质好文第一时间送达^_^

Click on the card below to follow US

>>>戳我一下,加入智驾机器人学习交流群✨

在三维视觉、SLAM、机器人定位里,Bundle Adjustment,简称 BA,几乎是绕不开的一步。

它干的事情可以说很朴素:给定一堆相机位姿、三维点和图像观测,反复调整相机和三维点,让三维点投影到图像上时尽可能对齐真实观测。说得更直白一点,BA 就是在“把相机摆正,把地图拧准”。

但这个老问题在今天遇到了一个新矛盾。

一边是传统 BA 工具链:Ceres、GTSAM、g2o 这些 C++ 库非常经典,工程上也被广泛使用。另一边是现在越来越常见的深度学习式感知系统:特征匹配、视觉 SLAM、SfM、三维重建,很多前端都已经跑在 PyTorch 里。

于是问题来了:神经网络在 GPU 上跑,BA 却往往要绕回 C++/CPU;模型想端到端训练,优化器却不在同一个计算图里;研究者想像写 PyTorch 一样调试 BA,传统因子图接口又不够灵活。

这篇论文 Bundle Adjustment in the Eager Mode 试图解决的,就是这个缝隙。

它做了一件很有意思的事:把大规模 BA 这种典型的二阶稀疏非线性优化,做成了一个能在 PyTorch Eager Mode 下运行的高效库,并且开源在:

https://github.com/pypose/bae

换句话说,作者想让 BA 不再只是“写 C++、搭因子图、调用 solver”的传统流程,而是变成一个更接近 PyTorch 生态的组件:能动态执行、能自动求导、能上 GPU、能和深度模型接在一起。

为什么 BA 放进 PyTorch 里并不简单?

如果只看公式,BA 很容易被误解成一个普通的最小二乘问题。

给定相机位姿、三维点和二维像素观测,BA 最小化的是重投影误差。也就是每个三维点通过相机模型投影到图像平面后,和真实匹配点之间的距离。

但是工程实现远没有公式这么轻松。

真正的大规模 BA 里,变量数量非常大:可能有几百、几千个相机位姿,几十万甚至更多三维点,以及海量图像观测。更关键的是,它的雅可比矩阵非常稀疏。

每一个残差只和一个相机、一个三维点有关。它并不关心其他相机,也不关心其他三维点。

这带来一个很重要的事实:如果把雅可比矩阵当成普通稠密矩阵来算,内存会爆炸。

论文里给了一个很直观的例子:BAL 数据集里的 Ladybug 场景,如果用标准 PyTorch 自动求导去构造稠密雅可比,在 double precision 下大约需要 5.2 TB 内存;但真正非零的梯度块只需要约 125 MB

这不是优化一下代码就能解决的问题,而是计算结构必须重新设计。

传统 BA 库强在哪里,又卡在哪里?

Ceres、GTSAM、g2o 这类工具的核心优势,是成熟、稳定、精度高。它们通常使用静态因子图,把变量和约束提前组织好,再交给优化器求解。

这套设计在经典 SLAM 和 SfM 里非常有效,但它和 PyTorch 这种动态图开发体验并不天然契合。

PyTorch 的 Eager Mode 是“写到哪,执行到哪”。你可以在 Python 里写 if、for、while,可以随时 print 中间变量,也可以用调试器一步步看。研究原型、快速试错、接神经网络训练,都很方便。

而传统 BA 库更像是先把图搭完,再统一执行。它不太适合在运行时频繁改变计算结构,也很难自然接入 PyTorch 的自动求导和 GPU 工作流。

这就导致几个现实问题:

第一,开发和调试不够顺手。研究者常常要在 Python 和 C++ 两边来回切换。

第二,深度网络和 BA 之间容易出现数据搬运。网络输出在 GPU 上,BA 求解在 CPU 或 C++ 后端,数据来回拷贝,延迟很难忽略。

第三,想做端到端训练时,BA 不容易成为同一个动态图里的自然组成部分。

这篇工作的目标不是重新发明 BA 的数学,而是把 BA 放进一个更适合现代深度学习研究的执行环境里。

这篇论文的核心想法:在 Eager Mode 里识别稀疏结构

作者提出的关键机制叫 sparsity-aware AutoDiff,可以理解为“知道稀疏结构的自动求导”。

普通 PyTorch 自动求导默认并不知道 BA 的稀疏关系。它看到的是一堆张量运算,很难天然知道“这个残差只依赖第 i 个相机和第 j 个点”。

论文的做法是:在 Eager Mode 执行过程中,自动追踪张量操作,把数据流表示成一个有向无环图,也就是 DAG。

这里有一个很巧妙的区分:

索引操作负责确定稀疏位置。比如 poses[cidx] 和 points[pidx] 会告诉系统:第 k 个观测对应哪个相机、哪个三维点。因此,雅可比矩阵中哪些块是非零的,可以从这些索引关系里推出来。

算术操作负责计算块里的数值。比如位姿作用到三维点、相机投影、残差计算、鲁棒核等操作,决定每个非零雅可比块的具体数值。

这样一分工,事情就清楚了:索引决定“哪里有梯度”,算术决定“梯度是多少”。

更重要的是,这个过程是在 PyTorch Eager Mode 中动态完成的,不需要用户提前手动搭静态因子图。用户仍然可以按 Python/PyTorch 的方式写模型,后端负责把稀疏结构抽出来。

为什么它能快?稀疏张量和 GPU 算子是关键

只知道雅可比稀疏还不够。LM 这类二阶优化方法还需要很多稀疏线性代数操作,比如:

  • 构造 normal equation: JᵀJ
  • 计算右端项:JᵀR
  • 加阻尼项:λ · diag(JᵀJ)
  • 解稀疏线性系统

PyTorch 原生对一些稀疏操作支持有限,尤其是 BA 需要的 block sparse 结构。因此,作者实现了一组面向 GPU 的稀疏算子,并尽量把它们封装成 PyTorch 风格的接口。

例如,稀疏矩阵乘法仍然可以用类似 mat1 @ mat2 的方式调用。对用户来说,它看起来像普通 PyTorch 运算;对后端来说,它用的是针对稀疏结构优化过的 GPU 实现。

论文中特别提到两类求解器:

一种是 Cholesky 直接求解器,适合中小规模问题。它通过分解矩阵获得相对精确的解。

另一种是 PCG,也就是预条件共轭梯度法,更适合大规模问题。它不用显式做完整分解,内存压力更小,并且更容易利用 GPU 并行。

为了减少重复开销,作者还缓存了稀疏矩阵乘法里的符号搜索结果,以及 Cholesky 分解中的符号分解结构。因为 BA 的稀疏模式在 LM 迭代过程中通常不变,只是数值在变,所以这些结构没必要每轮都重新算。

用户看到的接口:像写 PyTorch 一样写 BA

论文里给了一个最小可运行例子。代码的核心并不复杂:把相机位姿和三维点包成支持稀疏雅可比追踪的参数,然后在 forward 里通过索引取出对应相机和点,计算重投影误差,再交给 LM 优化器。

比较关键的是两处设计:

第一,pp.Parameter(..., sjac=True) 用来开启稀疏雅可比追踪。

第二,@psjac 标记残差函数,让后端知道这个函数可以按 batch 独立计算,从而更高效地组装稀疏雅可比。

这类接口设计的意义在于,它把复杂的后端细节藏起来了。研究者不需要亲手维护因子图,也不用手写每个雅可比块的位置,只需要按观测索引写出残差函数。

这对做机器人学习、神经 SLAM、深度 SfM 的研究者很有吸引力。因为这些任务往往不是单独跑一次 BA,而是需要把 BA 嵌入到更复杂的学习系统里。

实验结果:速度提升很明显,精度基本保持

论文在 BAL、1DSfM 和 CO3D v2 上做了实验。

在传统 BA benchmark 上,作者把自己的方法和 GTSAM、g2o、Ceres 等经典库进行了比较。

在 BAL 数据集上,作者方法整体比 GTSAM、g2o、Ceres 更快,同时最终重投影误差保持在同一量级。具体来看,Cholesky 版本在一些场景里非常快,例如 Trafalgar 和 Dubrovnik;PCG 版本则更适合大规模问题。

在 1DSfM 数据集上,优势更加明显。论文报告显示,该方法整体速度相比 GTSAM、g2o、Ceres 分别达到约  36×、43×、40× 的提升,同时误差没有明显牺牲。

在全文整体 benchmark 中,作者总结其 GPU eager-mode BA 相比 GTSAM、g2o、Ceres 的平均加速分别为 18.5×、22×、23×。

显存占用:不是靠“堆显存”硬冲

很多 GPU 加速方法容易让人担心一个问题:是不是只是把压力转移到显存上?

论文专门报告了 GPU memory usage。以 BAL 场景为例:

  • Dubrovnik:峰值约 2.84 GiB,平均约 1.8 GiB
  • Ladybug:峰值约 1.54 GiB,平均约 1002.4 MiB
  • Trafalgar:峰值约 539.75 MiB,平均约 350.4 MiB

作者观察到,当参数规模超过 30k 后,显存占用大体呈线性增长。峰值高于平均值,主要来自稀疏线性代数操作里的临时 buffer。

这说明它并不是用不可控的显存换速度,而是在稀疏结构下实现了比较可预测的内存增长。对后续接深度模型也很重要,因为 GPU 上还要留空间给网络本身。

和 GPU BA 方法比:不仅要快,还要能接 PyTorch 工作流

论文还比较了 Ceres CUDA、DeepLM、Theseus 等 GPU 相关方法。

结果显示,作者方法在 BAL 和 1DSfM 上整体快于 DeepLM。论文报告中,相比 DeepLM,作者方法在 BAL 上运行时间少约 **56%**,在 1DSfM 上少约 **28%**。

这里最值得注意的不是单纯速度,而是接口和生态位置。

DeepLM、Theseus 等方法虽然和 PyTorch 有关系,但并不是真正面向 PyTorch Eager Mode 的通用 BA 组件。它们往往依赖定制化数据结构或专用 CUDA 实现,扩展到其他优化问题并不自然。

而这篇工作强调的是:BA 只是开始,核心能力是“在 Eager Mode 里做稀疏二阶优化”。

接入 VGGT:让深度 SfM 后端更顺

为了证明这个库不是只会跑标准数据集,作者把它接进了 VGGT 这类深度学习 SfM pipeline。

VGGT 可以快速预测相机内外参、深度图、三维点云和点对应关系。但原始流程中,后处理 BA 依赖 PyCOLMAP,容易受到 CPU 计算和 CPU-GPU 数据传输的限制。

作者用自己的 BA 替换后端后,让整个 SfM pipeline 更自然地留在 PyTorch Eager Mode 里运行。

在 CO3D v2 上,VGGT + Ours 达到和 VGGT + PyCOLMAP 一样的 AUC@30,也就是 90.0,但时间从约 1.8s 降到约 0.9s。论文中还提到,平均 BA 优化时间约 0.7s,相比 PyCOLMAP 快约 2.3×

这部分实验很关键。因为它说明这个方法并不只是“单独跑 BA benchmark 很快”,而是可以嵌入现代深度三维视觉系统里,减少工程缝合成本。

接入 iMatching:BA 也能服务自监督特征匹配

除了 SfM 后处理,作者还把框架接入了 iMatching 这类自监督特征匹配方法。

iMatching 的思路是用 BA 作为自监督信号,让特征匹配模型通过几何一致性学习更好的对应关系。原始实现依赖 GTSAM,这意味着训练时需要在 PyTorch 和外部 C++ BA 后端之间做衔接。

作者将 BA 组件替换为自己的 eager-mode BA 后,梯度和优化流程可以更自然地留在 PyTorch 计算图里。

实验中,作者在 KITTI360 和 ETH3D-SLAM 上测试了特征匹配效果。在 KITTI360 上,iASpan 和 iDKM 的表现与原方法一致或更好;在 ETH3D-SLAM 上,整体结果也保持稳定,说明替换 BA 后端不会破坏原有学习框架,甚至能简化实现。

不只 BA:还能推广到 PGO

论文最后还做了一个很重要的扩展:Pose Graph Optimization,简称 PGO。

PGO 和 BA 一样,也是机器人和 SLAM 里常见的稀疏非线性最小二乘问题。不同的是,BA 同时优化相机和三维点;PGO 主要优化一组位姿,约束来自相对位姿测量,比如里程计和回环。

作者的方法可以自然迁移到 PGO,因为核心机制仍然成立:

索引操作决定哪个残差连接哪两个位姿,Lie group 运算决定残差和雅可比块的数值,后端稀疏 LM 负责求解。

在 parking-garage 和 sphere-a 两个 PGO benchmark 上,作者方法相比没有稀疏支持的 PyPose 快很多。论文中报告,在 PGO 上作者方法比 PyPose 快约 659×。在 sphere-a 这种较大样本上,Cholesky 版本相比 Ceres 也有明显速度优势。

这说明这篇工作的价值不局限于 BA。它更像是在 PyTorch 里补上了一块“稀疏二阶优化基础设施”。

这篇工作的意义:让几何优化重新靠近深度学习系统

过去几年,三维视觉领域有一个明显趋势:前端越来越神经网络化,但后端几何优化仍然很经典。

这并不是坏事。几何优化的可解释性、精度和稳定性仍然非常重要。但如果前端和后端长期处在两套生态里,研究和工程都会被割裂。

这篇论文的贡献,恰好是在两者之间搭桥:

它没有放弃 BA 的稀疏二阶优化结构,也没有强行把问题变成普通一阶梯度下降;同时,它又让 BA 能以 PyTorch Eager Mode 的方式存在,能调试、能接模型、能上 GPU、能扩展到 PGO。

可以把它概括成三点:

第一,作者提出了一个真正面向 PyTorch Eager Mode 的 BA 库,让传统依赖静态因子图的优化过程可以在运行时动态图里完成。

第二,作者设计了 sparsity-aware AutoDiff,通过 DAG 追踪索引和算术操作,自动推断稀疏雅可比结构,并计算非零雅可比块。

第三,作者补齐了 GPU 稀疏线性代数和求解器,包括稀疏矩阵乘法、对角阻尼、Cholesky 和 PCG,使 LM 这类二阶方法能在 PyTorch 生态中高效运行。

对做机器人、SLAM、SfM、神经渲染、三维重建的人来说,这类工具最大的价值可能不是“又快了多少倍”,而是它降低了把几何优化塞进学习系统里的门槛。

以前要跨语言、跨设备、跨计算图。现在,至少有了一条更 PyTorch、更动态、更容易实验的路。

也有一些值得继续看的问题

当然,这篇工作也不是终点。

作者在论文结尾提到,当前实现更偏向 GPU。未来如果加入更强的 CPU 多线程、SIMD 优化,也许能让 CPU 场景更好用。

另外,Python 自动内存管理和 PyTorch 动态张量管理会带来额外内存开销。后续如果结合手动 buffer 预分配、TorchDynamo 或 TorchScript,也可能进一步压低运行时间和显存使用。

再往后看,这套 eager-mode 稀疏优化机制也许可以扩展到更多机器人问题,比如机械臂操作、轨迹优化、运动规划,甚至更多类型的非线性优化器,比如 dogleg、L-BFGS 等。

如果这些方向继续完善,PyTorch 里的几何优化可能会变得越来越像一个“标准组件”,而不是一个需要额外接驳的外部系统。

论文信息


论文标题: Bundle Adjustment in the Eager Mode

作者: Zitong Zhan, Huan Xu, Zihang Fang, Xinpeng Wei, Yaoyu Hu, Chen Wang

机构: University at Buffalo;Georgia Institute of Technology;Purdue University;Carnegie Mellon University

开源项目: https://github.com/pypose/bae

| 本文仅做学术分享,如有侵权,请联系删文!

点赞
转发
推荐
评论

点击阅读原文查看更多内容

往期推荐


VIO 初始化能不能不追特征?


机器人定位不用“两张地图”了?SOCC-ICP 让语义占据栅格直接接管 LiDAR 里程计


顶尖学术助力
自动驾驶与机器人领域论文/专利全程辅导
01
全流程论文辅导

(1)自动驾驶可指导方CUDA编程高性能计算HPCCV/感知算法端到端自动驾驶决策规划显著性分析图像分割LLM自动驾驶雷达感知自动驾驶感知毫米波雷达深度学习滤波算法预期功能安全自动驾驶基础共性技术研究自动驾驶模拟仿真技术研究自动驾驶安全性设计及验证仿真与测试场景生成强化学习,预期功能安全,自动驾驶点云处理,行为识别,目标检测,视觉感知,BEV感知,边缘计算,数据处理,驾驶行为研究,点云,多模态,自动驾驶决策规划,英伟达平台模型部署优化,自动驾驶安全方向等。

(2)机器人可指导方向:机器人路径规划及算法AI 集中在ROS机器人和CV NLP计算机视觉,机器学习,机器人三维视觉, 图像融合 ,图像理解 ,机器人算法, SLAM,点云处理,信号处理,具身智能,智能控制,机器人柔顺控制,分数阶控制,自适应反步,产业机器人,电力检测机器人,海洋机器人,进化计算,移动机器人定位导航,位姿估计,轮式机器人,仿生足式机器人机器人感知,语义分割,深度学习,机器视觉,工业机器人移动机器人,机器人模仿学习,控制算法设计,多模态智能等。

02
专利布局与成果转化


我们专注自动驾驶和机器人等前沿领域,提供选题创新性评估、实验设计、论文写作与顶级会议/期刊投稿指导。团队源自全球顶尖实验室及企业研究院,强化创新点与工程实现,不仅保证论文产出,更传授科研思维与方法,助你掌握独立发表高水平论文的能力。提供专利挖掘、技术交底书撰写、国内外专利申请(发明专利/实用新型)全流程服务。结合产业需求,强化权利要求的保护范围与商业价值,助力成果转化与竞争力提升。


为什么选择我们?
▪️ 导师团队有3000+QS前50博导库100% 博士团队
▪️ 可指导期刊or会议:ICCV、IROS、ICRA、CVPR、AAAI、ICLR、T-PAMI、IJCV、T-RO、RAL、T-IV、TIP、TIM、TMECH、TVT、TASE、IEEE IoT J.  等,并提供博士硕士毕业论文全程辅导及专利授权服务。
▪️ 大牛阵容:Nature作者、顶级期刊/会议审稿人、研究员、哈耶普斯麻牛剑导师带队定制一对一学习课纲
图片
扫码预约大咖导师一对一指导↑↑↑

Python社区是高质量的Python/Django开发社区
本文地址:http://www.python88.com/topic/198034