把 Bundle Adjustment 塞进 PyTorch：BAE 为什么不只是“用 autograd 写个残差”

0. 简介捆绑调整Bundle Adjustment, BA在 3D vision、SLAM、SfM 里一直是绕不过去的模块。它做的事看起来很朴素给定多张图像里的 2D 观测、相机位姿、相机内参和 3D 点反复调整相机和点的位置让 3D 点重新投影到图像上以后尽量贴近观测像素。问题在于BA 不是普通的深度学习 loss它通常需要二阶优化、稀疏雅可比矩阵、李群上的位姿更新和专门的稀疏线性求解器。传统做法是把这部分交给 Ceres、g2o、GTSAM 这类 C 优化库神经网络在 PyTorch 里跑几何优化在另一个世界里跑。这篇 Bundle Adjustment in the Eager Mode 的核心判断很直接BA 不应该永远停留在 PyTorch 外面。论文和bae项目把传统依赖静态 factor graph 的 BA搬进了 PyTorch 的即时执行模式Eager Mode并且不是玩具级演示。按照论文报告在 BAL 和 1DSfM 等数据集上它在 GPU 上相对 GTSAM、g2o、Ceres 分别取得约 18.5 倍、22 倍、23 倍的平均加速在部分大规模场景里相对传统 CPU 优化器可以接近甚至超过百倍速度提升。这里的关键不是“Python 比 C 快”而是 BA 的稀疏结构被正确识别后GPU 并行度终于能被用起来。进一步看这个工作最有工程价值的地方是它没有要求用户重新学习一套陌生的 factor graph DSL。用户仍然写nn.Module仍然用 tensor indexing 表达哪个 observation 属于哪个 camera、哪个 point仍然用optimizer.step(...)跑优化。差别只在几个关键位置用pp.Parameter(..., sjacTrue)告诉系统这些参数需要稀疏雅可比追踪用psjac标记可以并行装配稀疏 Jacobian 的残差函数再用LM(..., sparseTrue)或bae.optim.LM把问题交给稀疏 LM 后端。换句话说它试图保留 PyTorch 的开发体验同时补上 PyTorch 原本缺失的稀疏二阶优化能力。1. BA 到底在优化什么1.1 从重投影误差开始如果你第一次接触 BA可以先把它理解成一个“多相机、多点的全局对齐问题”。假设第i ii个相机看到第j jj个 3D 点图像上的观测像素是x i j x_{ij}xij​当前估计的相机位姿是ζ i \zeta_iζi​3D 点坐标是p j \mathbf{p}_jpj​相机投影模型是Π \PiΠ那么 BA 要最小化的就是投影点和观测点之间的差。所有被观测到的点都参与求和最终形成一个非线性最小二乘问题。真实场景里不是每个相机都能看到每个点所以观测关系天然是稀疏的这也是后面所有工程设计的出发点。θ ∗ arg ⁡ min ⁡ θ ∑ ( i , j ) ∈ O ∥ Π ( ζ i , p j , K i ) − x i j ∥ 2 2 \theta^{*}\arg\min_{\theta}\sum_{(i,j)\in\mathcal{O}} \left\|\Pi(\zeta_i,\mathbf{p}_j,\mathbf{K}_i)-\mathbf{x}_{ij}\right\|_2^2θ∗argθmin​(i,j)∈O∑​∥Π(ζi​,pj​,Ki​)−xij​∥22​这里的θ \thetaθ通常由两大类变量拼起来第一类是相机变量包括位姿、焦距、畸变等第二类是三维路标点也就是场景里的 3D landmark。论文里相机位姿用S E ( 3 ) SE(3)SE(3)表示外部存储常见是 7 维形式例如平移加四元数但优化更新发生在 6 维切空间里。这一点很容易被忽略你看到代码里有camera_params[..., :7]并不意味着优化器真的把四元数当作普通 7 维欧式向量乱加它背后要处理李群和李代数之间的映射。1.2 LM 为什么比普通 Adam 更合适BA 通常不是拿 Adam 或 SGD 硬推因为重投影误差的局部几何结构非常清楚使用二阶近似更高效。论文采用的是列文伯格-马夸尔特Levenberg-Marquardt, LM它可以看成 Gauss-Newton 和梯度下降之间的折中当局部模型可信时它像 Gauss-Newton 一样快速下降当问题病态或步子太大时阻尼项会让更新更保守。对 SLAM 和 SfM 这种强几何问题来说LM 的稳定性和收敛速度往往比一阶优化器更合适。( J ⊤ J λ ⋅ d i a g ( J ⊤ J ) ) Δ θ − J ⊤ R \left(\mathbf{J}^{\top}\mathbf{J}\lambda\cdot\mathrm{diag}(\mathbf{J}^{\top}\mathbf{J})\right)\Delta\theta -\mathbf{J}^{\top}\mathbf{R}(J⊤Jλ⋅diag(J⊤J))Δθ−J⊤Rθ ← θ Δ θ \theta \leftarrow \theta\Delta\thetaθ←θΔθ这里R \mathbf{R}R是所有 residual 拼起来的向量J \mathbf{J}J是 residual 对参数的 Jacobianλ \lambdaλ是阻尼系数。真正贵的地方不在公式本身而在J \mathbf{J}J的构建、J ⊤ J \mathbf{J}^{\top}\mathbf{J}J⊤J的稀疏乘法、对角阻尼、以及稀疏线性系统求解。传统 C solver 的优势正是在这些地方它们知道 factor graph 的结构知道每个 residual 只连到少量变量所以从一开始就不会把问题当成一个巨大的 dense matrix。2. 为什么“直接用 PyTorch autograd”不够2.1 核心问题在于 PyTorch 不知道 BA 的稀疏性把 BA 塞进 PyTorch最自然的第一反应是写一个 residual function然后让 autograd 自动求导不就行了吗问题就在这里。标准 PyTorch AutoDiff 可以求导但它并不知道 BA 的 block sparsity。对 PyTorch 来说输出 residual 和输入参数都是大 tensor如果没有额外结构提示它倾向于考虑“所有输出对所有参数”的梯度关系。可是 BA 的真实依赖关系很窄一个 2D residual 只依赖一个相机和一个 3D 点和其他相机、其他点没有任何关系。论文给了一个很直观的数字BAL 数据集里的 Ladybug 场景包含 1,723 个相机和 156k 个 3D 点如果使用 dense Jacobian双精度下内存会到 5.2 TB但只存真实非零块稀疏 Jacobian 大约 125 MB。这个差异不是小优化而是能不能运行的问题。很多人说“PyTorch 做 BA 慢”本质上不是 Python 循环慢这么简单而是没有把 BA 的稀疏数学结构表达给底层系统。2.2 indexing 决定结构算术运算决定数值BAE 的一个关键拆分是tensor indexing 决定 Jacobian 的稀疏结构投影、李群变换、畸变模型这些算术运算决定非零 block 的具体数值。比如poses[cidx]表示第k个 observation 用第cidx[k]个相机points[pidx]表示第k个 observation 用第pidx[k]个 3D 点。这两个 index tensor 已经把 residual 和变量之间的连接关系说清楚了后端只需要沿着执行时的数据流把这些关系记录下来。换句话说用户没有手写 factor graph但代码本身已经隐含了 factor graph。BAE 做的事是在 eager mode 运行时追踪这些 tensor 操作把数据流组织成一个有向无环图然后在反向阶段根据索引关系把 Jacobian block 放回正确位置。算术部分则用 batched Jacobian 计算非零块的值不需要对每个 block 单独 backward。这里的设计很有意思它不是绕开 PyTorch而是把 PyTorch 的动态执行方式变成稀疏结构推断的输入。3. 最小 BA 代码看懂sjac和psjac3.1 一个最小可跑的 PyPose 风格例子下面这段代码保留了论文里的最小示意。它没有加载真实 BAL 数据集只构造一个相机和 8 个 3D 点用随机像素观测演示 API 形态。真实项目里你会把pixels、cidx、pidx换成数据集或前端匹配模块输出的观测把相机和点初始化换成 SfM、SLAM 或神经网络预测的结果。需要注意的是pp.Parameter(..., sjacTrue)和psjac不是装饰语法它们分别打开了稀疏 Jacobian 追踪和批量装配路径。importtorchimportpyposeasppfromtorchimportnnfrompypose.optimimportLMfrompypose.optim.solverimportPCGfrompypose.optim.strategyimportTrustRegionfrompypose.optim.schedulerimportStopOnPlateaufrompypose.autograd.functionimportpsjacclassReprojErr(nn.Module):def__init__(self,poses,points):super().__init__()self.posespp.Parameter(poses,sjacTrue)self.pointspp.Parameter(points,sjacTrue)psjacdefproject(poses,points):pointsposes.Act(points)return-points[...,:2]/points[...,[2]]defforward(self,pixels,cidx,pidx):posesself.poses[cidx]pointsself.points[pidx]returnReprojErr.project(poses,points)-pixels torch.set_default_device(cuda)npts,poses8,pp.randn_SE3(1)pointstorch.randn(npts,3)points[:,2]4cidxtorch.zeros(npts,dtypetorch.long)pidxtorch.arange(npts)pixelstorch.randn(npts,2)inputs(pixels,cidx,pidx)modelReprojErr(poses,points)solverPCG(tol1e-4,maxiter250)strategyTrustRegion(up2.0,down0.5**4)optimizerLM(model,solver,strategy,sparseTrue)schedulerStopOnPlateau(optimizer,steps5,verboseTrue)whilescheduler.continual():lossoptimizer.step(inputs)scheduler.step(loss)这段代码里最值得盯住的是forward。self.poses[cidx]和self.points[pidx]会把每条 observation 对应的相机和点复制出来使后面的投影可以批量计算。直观理解是原来是“很多 residual 共享相机和点”现在在计算布局上变成“每条 residual 有自己对齐的一份相机和点引用”。这个复制不是为了改变数学问题而是为了让 GPU 能用连续批量操作高效计算同时让稀疏 Jacobian 后端知道每个 residual 的非零 block 应该放在哪一行、哪一列。3.2 真实 BAL 场景里还会优化内参和畸变官方bae仓库里的 BAL 示例比最小例子更接近真实 BA。相机参数通常包含SE(3)位姿、焦距f、径向畸变k1和k2。代码里camera_params[..., :7]被解释成SE(3)后 3 维参与投影模型。这个设计也解释了为什么论文和代码里会同时出现 7 维位姿表示、6 维梯度空间、10 维相机参数这几种说法存储、几何更新和实际优化变量数不是同一个层面。psjacdefproject(points,camera_params):projectionpp.SE3(camera_params[...,:7]).Act(points)projection-projection[...,:2]/projection[...,[2]]fcamera_params[...,[-3]]k1camera_params[...,[-2]]k2camera_params[...,[-1]]ntorch.sum(projection**2,axis-1,keepdimTrue)r1k1*nk2*n**2returnprojection*r*fclassResidual(nn.Module):def__init__(self,camera_params,points):super().__init__()self.posepp.Parameter(camera_params,sjacTrue)self.pointspp.Parameter(points,sjacTrue)self.pose.trim_SE3_gradTruedefforward(self,observes,cidx,pidx):points_projproject(self.points[pidx],self.pose[cidx])returnpoints_proj-observes这里的trim_SE3_grad True也值得单独看。SE(3)位姿可以用 7 个数存但真正的李代数更新只有 6 个自由度所以优化器需要把冗余表示和有效梯度空间区分开。对读者来说这段代码的意义不在于畸变模型有多复杂而在于它说明 BAE 不是只能跑“无内参、无畸变、无李群”的教学例子。它保留了 PyTorch module 的外观但内部仍然处理 BA 里真实存在的几何细节。4. 稀疏 Jacobian 怎么存为什么选择 PyTorch sparse BSR4.1 BSR 比 COO 更贴近 BA 的 block 结构BA 的 Jacobian 不是随机稀疏而是 block sparse。每条 2D residual 对相机通常产生一个2 × 6 2\times62×6或更宽的 block对 3D 点产生一个2 × 3 2\times32×3block。用普通坐标列表格式当然也能存但它会把 block 拆成一个个标量非零项很多结构信息就丢了。BAE 采用 PyTorch 原生稀疏张量尤其是分块稀疏行格式Block Sparse Row, BSR让 Jacobian 仍然是 tensor只是 layout 是 sparse BSR。这个选择有两层含义。第一它让 BAE 不需要发明一套和 PyTorch 脱节的 Jacobian dictionary用户代码和优化器接口都更像 PyTorch。第二BSR 对J ⊤ J \mathbf{J}^{\top}\mathbf{J}J⊤J这类 block 稀疏矩阵乘法更自然因为乘法的最小单位可以是 block而不是单个浮点数。PyTorch 官方文档本身就提供了 sparse tensor 体系包括 COO、CSR、CSC、BSR、BSC 等格式BAE 的问题在于PyTorch 对 BSR 上很多二阶优化需要的操作支持并不完整所以它必须补一批算子。4.2 不是只有存储关键还在稀疏线性代数LM 每一步都绕不开J ⊤ J \mathbf{J}^{\top}\mathbf{J}J⊤J、J ⊤ R \mathbf{J}^{\top}\mathbf{R}J⊤R、对角阻尼、线性系统求解。论文里把这些操作拆成了几类稀疏矩阵-矩阵乘法、稀疏矩阵-向量乘法、对角元素 clamp/scale、稀疏 Cholesky 或 PCG 求解。PyTorch 原生支持其中一部分但不够覆盖完整 BA 后端。BAE 因此实现了 GPU sparse operators并注册到 PyTorch operator dispatcher 里让用户依旧可以写接近普通 tensor 的表达式例如A J.T J。进一步看J ⊤ J \mathbf{J}^{\top}\mathbf{J}J⊤J的 sparse matrix multiplication 通常分成 symbolic 和 numeric 两个阶段。symbolic 阶段只看稀疏结构决定输出矩阵哪些 block 会非零numeric 阶段才真正做 block 乘加。BA 的 sparsity pattern 在 LM 迭代过程中通常不变因为哪台相机看到哪个点不会在优化中频繁改变所以 symbolic 结果可以缓存复用。这个细节很工程但正是这类工程细节决定了 eager mode BA 能不能从“能跑”变成“跑得快”。5. 求解器PCG、Cholesky 和 Schur complement 的取舍5.1 PCG 适合大规模Cholesky 适合中小规模或特定结构BAE 支持不同线性求解策略。论文实验里出现了 PCG 和 Cholesky 两类路径。**预条件共轭梯度Preconditioned Conjugate Gradient, PCG**是迭代法每轮主要做 sparse matrix-vector product 和向量更新内存占用相对可控更适合大规模系统。Cholesky 是直接法对中小规模问题或某些结构更快但分解过程可能需要更多工作空间。官方仓库里的PCG还会从矩阵对角线构造一个简单预条件器缓解线性系统条件数带来的收敛问题。classPCG(CG):def__init__(self,maxiterNone,tol1e-5):super().__init__(maxiter,tol)defforward(self,A,b,xNone,MNone):l_diagA.diagonal()l_diag[l_diag.abs()1e-6]1e-6Mspdiags_((1/l_diag),None,shapeA.shape,layoutNone)ifA.layouttorch.sparse_bsr:MM.to_sparse_bsr(blocksizeA.values().shape[-2:]).to(A.device)returnsuper().forward(A,b,x,M)这段简化后的代码说明了一个常见工程判断大规模 BA 不一定非要把矩阵完整分解很多时候只要能高效地做A x迭代求解就足够有竞争力。论文附录还提到 PCG 可以用 CUDA Graph capture/replay 减少重复 kernel launch 的开销。这个思路很符合 PyTorch eager mode 的现实Python 写法可读但反复启动小 kernel 有成本把重复迭代捕获成 CUDA Graph可以在不牺牲接口可读性的前提下接近更底层的执行效率。5.2 Schur complement 是 BA 的老朋友BA 还有一个经典结构变量可以分成相机和 3D 点。如果把 normal equation 按这两类变量分块会得到如下结构。由于每个 3D 点只和看见它的观测相关点块V \mathbf{V}V往往是 block diagonal求逆或应用逆都比较便宜。因此可以先消去点变量只求一个规模更小的相机系统这就是舒尔补Schur Complement。传统 BA solver 很重视这一步因为相机数量通常远小于点数量。[ U W W ⊤ V ] [ Δ c Δ p ] [ g c g p ] \begin{bmatrix} \mathbf{U} \mathbf{W}\\ \mathbf{W}^{\top} \mathbf{V} \end{bmatrix} \begin{bmatrix} \Delta\mathbf{c}\\ \Delta\mathbf{p} \end{bmatrix} \begin{bmatrix} \mathbf{g}_c\\ \mathbf{g}_p \end{bmatrix}[UW⊤​WV​][ΔcΔp​][gc​gp​​]S Δ c g c − W V − 1 g p , S U − W V − 1 W ⊤ \mathbf{S}\Delta\mathbf{c}\mathbf{g}_c-\mathbf{W}\mathbf{V}^{-1}\mathbf{g}_p,\quad \mathbf{S}\mathbf{U}-\mathbf{W}\mathbf{V}^{-1}\mathbf{W}^{\top}SΔcgc​−WV−1gp​,SU−WV−1W⊤Δ p V − 1 ( g p − W ⊤ Δ c ) \Delta\mathbf{p}\mathbf{V}^{-1}\left(\mathbf{g}_p-\mathbf{W}^{\top}\Delta\mathbf{c}\right)ΔpV−1(gp​−W⊤Δc)项目当中还有一组关于 Schur 和 matrix-free 的内存对比在最大的venice/problem-1778这类约 300 万可优化参数问题上Schur matrix-free 的峰值显存约 5.8 GiB平均约 2.3 GiB相比 LM matrix-free峰值显存少约 43%平均显存少约 61%。这说明 BAE 后续已经不只是复现论文里的稀疏 LM还在继续沿着传统 BA 的核心工程路线做优化消元、matrix-free、缓存结构、减少显存占用。…详情请参照古月居