bae-kai 0.1.1

PyTorch implementation of Bundle Adjustment in the Eager-mode

bae-kai: Bundle Adjustment in the Eager-mode

⚠️ Development Phase Notice: This library is currently in active development. APIs are subject to change and should be considered experimental. Use at your own discretion in production environments.

bae-kai is a fork of bae with full Windows support, pre-built CUDA wheels, and cuDSS bundling.

bae is a PyTorch-based library supporting 2nd-order optimization techniques. The library provides efficient implementations for sparse optimization problems in robotics, particularly Bundle Adjustment (BA) and Pose Graph Optimization (PGO).

News

• 2026-02-06: Windows CUDA wheels now available on PyPI
• 2025-12-12: Added a VGGT integration example.

Features

• Sparse Block Matrix Operations: Optimized implementations of sparse matrix operations for large-scale optimization
• CUDA Acceleration: Custom CUDA kernels for high-performance sparse linear algebra
• Bundle Adjustment: Efficient implementation for camera pose and 3D structure optimization
• Pose Graph Optimization: Tools for optimizing robot trajectories using pose graph representations
• PyTorch Integration: Seamlessly integrates with PyTorch's automatic differentiation framework
• Levenberg-Marquardt Optimizer: Custom implementation of the LM algorithm for non-linear least squares problems

Future Plan

• An new backend for distributed solver

Installation

Quick Install (Recommended)

Pre-built wheels with CUDA extensions are available on PyPI:

# Install with CUDA 12.8 (for RTX 30/40/50 series)
uv pip install torch --index-url https://download.pytorch.org/whl/cu128
uv pip install bae-kai

# Or with CUDA 13.0 (latest)
uv pip install torch --index-url https://download.pytorch.org/whl/cu130
uv pip install bae-kai

Available Wheels

Platform	CUDA	Architectures
Linux	12.4, 12.8, 13.0	sm_70 - sm_120
Windows	12.4, 12.6, 12.8	sm_70 - sm_120

From Source

For development or custom builds:

# Clone this repository
git clone https://github.com/OpsiClear/bae-kai.git
cd bae-kai

# Install PyPose from the bae branch
uv pip install git+https://github.com/pypose/pypose.git@bae

# Install in development mode (uv sync handles all dependencies)
uv sync

Build Options

Control the build with environment variables:

• BAE_BUILD_EXTENSIONS=1: Force building CUDA extensions (required on Windows)
• BAE_SKIP_EXTENSIONS=1: Skip CUDA extensions (Python fallbacks not available)
• USE_CUDSS: Set to "1" (default) to enable cuDSS support, "0" to disable
• CUDSS_DIR: Path to cuDSS installation if not in standard locations

Example Usage

Bundle Adjustment

Bundle Adjustment optimizes camera poses and 3D point positions to minimize reprojection error. The following example shows how to perform BA using bae:

import torch
from datapipes.bal_loader import get_problem
from ba_helpers import Reproj
from bae.optim import LM

# Load a problem from the BAL dataset
dataset = get_problem("problem-49-7776-pre", "ladybug", use_quat=True)
dataset = {k: v.to('cuda') for k, v in dataset.items() if isinstance(v, torch.Tensor)}

# Prepare input for the optimization
input = {
    "points_2d": dataset['points_2d'],
    "camera_indices": dataset['camera_index_of_observations'],
    "point_indices": dataset['point_index_of_observations']
}

# Initialize model with camera parameters and 3D points
model = Reproj(
    dataset['camera_params'].clone(),
    dataset['points_3d'].clone()
).to('cuda')

# Auto-selection: solver, strategy, and method are chosen automatically
optimizer = LM(model, reject=30)

# Run optimization for multiple iterations
for idx in range(20):
    loss = optimizer.step(input)
    print(f'Iteration {idx}, loss: {loss.item()}')

For more control, you can configure the optimizer explicitly:

# String-based configuration
optimizer = LM(model, solver="pcg", strategy="trustregion", method="schur")

# Object-based configuration (backward-compatible)
from bae.utils.pysolvers import PCG
from bae.utils.schur import TrustRegion
optimizer = LM(model, solver=PCG(tol=1e-4, maxiter=250), strategy=TrustRegion())

See ba_example.py for a complete working example.

API Overview

Module	Exports	Description
bae.optim	LM, SchurLM	Levenberg-Marquardt optimizer with auto-selection
bae.autograd	TrackingTensor, map_transform, jacobian	Sparse jacobian via operation tracing
bae.utils	PCG, PCG_, CuDSS, SciPySpSolver	Linear solvers
bae.utils	TrustRegion, Adaptive	Damping strategies

Integration with VGGT

bae is used as an optional Bundle Adjustment backend in our VGGT fork (Visual Geometry Grounded Transformer) to refine the camera poses, intrinsics, and 3D points predicted by VGGT before exporting a COLMAP reconstruction.

After installing bae, you can run VGGT's COLMAP export with BA enabled and bae selected as the solver:

python demo_colmap.py --scene_dir /path/to/scene --use_ba --implementation bae  # optional: --shared_camera

This command invokes prepare_bae(...) inside vggt/demo_colmap.py, which wraps VGGT tracks and predictions into bae.optim.LM and updates extrinsic, intrinsic, and points_3d in place before writing scene_dir/sparse/ in COLMAP format.

Dataset Support

The library supports common optimization datasets and tasks:

• Bundle Adjustment in the Large (BAL) dataset
• 1DSfM dataset for large-scale structure from motion
• G2O pose graph datasets

Performance

bae is designed for high performance using:

• Efficient sparse block matrix operations
• CUDA acceleration for core operations
• Optimized linear solvers (PCG, CUDA Sparse Solver)
• Memory-efficient data structures

Citation

If you use bae in your research, please cite:

@article{zhan2025bundle,
  title = {Bundle Adjustment in the Eager Mode},
  author = {Zhan, Zitong and Xu, Huan and Fang, Zihang and Wei, Xinpeng and Hu, Yaoyu and Wang, Chen},
  journal = {arXiv preprint arXiv:2409.12190},
  year = {2025},
  url = {https://arxiv.org/abs/2409.12190}
}

Acknowledgements

This project is a fork of bae by Zitong Zhan et al.

The implementation draws inspiration from:

• bae (original) - Bundle Adjustment in the Eager Mode
• PyPose for SE(3) pose representations
• GTSAM for reprojection jacobian concepts

License

This project is licensed under the GNU Affero General Public License v3.0.