From 71f8a4ced3c167de0ffe911caf1dd2278d23a836 Mon Sep 17 00:00:00 2001 From: brentyi Date: Fri, 6 Dec 2024 12:18:04 +0000 Subject: [PATCH] Update README, fix pose graph timing --- README.md | 95 ++++++++++++++++++-------------------- examples/pose_graph_g2o.py | 7 +++ 2 files changed, 51 insertions(+), 51 deletions(-) diff --git a/README.md b/README.md index 1a4356a..5feb0de 100644 --- a/README.md +++ b/README.md @@ -2,17 +2,23 @@ [![pyright](https://github.com/brentyi/jaxls/actions/workflows/pyright.yml/badge.svg)](https://github.com/brentyi/jaxls/actions/workflows/pyright.yml) -**`jaxls`** is a library for nonlinear least squares in JAX. +**`jaxls`** is a library for solving sparse [NLLS](https://en.wikipedia.org/wiki/Non-linear_least_squares) and [IRLS](https://en.wikipedia.org/wiki/Iteratively_reweighted_least_squares) problems in JAX. +These are common in classical robotics and computer vision, as well as for MAP inference with Gaussian likelihood functions. + +To install on Python 3.12+: + +```bash +pip install git+https://github.com/brentyi/jaxls.git +``` + +### Overviews We provide a factor graph interface for specifying and solving least squares -problems. We accelerate optimization by analyzing the structure of graphs: +problems. **`jaxls`** takes adantage of structure in graphs: repeated factor and variable types are vectorized, and the sparsity of adjacency -in the graph is translated into sparse matrix operations. - -Use cases are primarily in least squares problems that are (1) sparse and (2) -inefficient to solve with gradient-based methods. +is translated into sparse matrix operations. -Currently supported: +Supported: - Automatic sparse Jacobians. - Optimization on manifolds. @@ -22,46 +28,16 @@ Currently supported: - Sparse iterative with Conjugate Gradient. - Preconditioning: block and point Jacobi. - Inexact Newton via Eisenstat-Walker. - - Sparse direct with Cholesky / CHOLMOD, on CPU. - - Dense Cholesky for smaller problems. + - Recommended for most problems. + - Dense Cholesky. + - Fast for small problems. + - Sparse Cholesky, on CPU. (CHOLMOD) -For the first iteration of this library, written for [IROS 2021](https://github.com/brentyi/dfgo), see -[jaxfg](https://github.com/brentyi/jaxfg). `jaxls` is a rewrite that is faster -and easier to use. For additional references, see inspirations like +`jaxls` borrows heavily from libraries like [GTSAM](https://gtsam.org/), [Ceres Solver](http://ceres-solver.org/), [minisam](https://github.com/dongjing3309/minisam), [SwiftFusion](https://github.com/borglab/SwiftFusion), -[g2o](https://github.com/RainerKuemmerle/g2o). - -### Installation - -`jaxls` supports `python>=3.12`: - -```bash -pip install git+https://github.com/brentyi/jaxls.git -``` - -**Optional: CHOLMOD dependencies** - -By default, we use an iterative linear solver. This requires no extra -dependencies. For some problems, like those with banded matrices, a direct -solver can be much faster. - -For Cholesky factorization via CHOLMOD, we rely on SuiteSparse: - -```bash -# Option 1: via conda. -conda install conda-forge::suitesparse - -# Option 2: via apt. -sudo apt install -y libsuitesparse-dev -``` - -You'll also need _scikit-sparse_: - -```bash -pip install scikit-sparse -``` +and [g2o](https://github.com/RainerKuemmerle/g2o). ### Pose graph example @@ -111,8 +87,12 @@ factors = [ Factors with similar structure, like the first two in this example, will be vectorized under-the-hood. -**Solving optimization problems.** We can set up the optimization problem, solve -it, and print the solutions: +Batched inputs can also be manually constructed, and are detected by inspecting +the shape of variable ID arrays in the input. Just add a leading batch axis to +all elements in the arguments tuple. + +**Solving optimization problems.** To set up the optimization problem, solve +it, and print solutions: ```python graph = jaxls.FactorGraph.make(factors, pose_vars) @@ -122,11 +102,24 @@ print("Pose 0", solution[pose_vars[0]]) print("Pose 1", solution[pose_vars[1]]) ``` -### Limitations +### CHOLMOD setup + +By default, we use an iterative linear solver. This requires no extra +dependencies. For some problems, like those with banded matrices, a direct +solver can be much faster. + +For Cholesky factorization via CHOLMOD, we rely on SuiteSparse: -There are many practical features that we don't currently support: +```bash +# Option 1: via conda. +conda install conda-forge::suitesparse -- GPU accelerated Cholesky factorization. (for CHOLMOD we wrap [scikit-sparse](https://scikit-sparse.readthedocs.io/en/latest/), which runs on CPU only) -- Covariance estimation / marginalization. -- Incremental solves. -- Analytical Jacobians. +# Option 2: via apt. +sudo apt install -y libsuitesparse-dev +``` + +You'll also need _scikit-sparse_: + +```bash +pip install scikit-sparse +``` diff --git a/examples/pose_graph_g2o.py b/examples/pose_graph_g2o.py index 4b2568d..cfa2010 100755 --- a/examples/pose_graph_g2o.py +++ b/examples/pose_graph_g2o.py @@ -16,6 +16,10 @@ import _g2o_utils +jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache") +jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1) +jax.config.update("jax_persistent_cache_min_compile_time_secs", 0) + def main( g2o_path: pathlib.Path = pathlib.Path(__file__).parent / "data/input_M3500_g2o.g2o", @@ -40,16 +44,19 @@ def main( for pose_var, pose in zip(g2o.pose_vars, g2o.initial_poses) ) ) + jax.block_until_ready(initial_vals) with jaxls.utils.stopwatch("Running solve"): solution_vals = graph.solve( initial_vals, trust_region=None, linear_solver=linear_solver ) + jax.block_until_ready(solution_vals) with jaxls.utils.stopwatch("Running solve (again)"): solution_vals = graph.solve( initial_vals, trust_region=None, linear_solver=linear_solver ) + jax.block_until_ready(solution_vals) # Plot plt.figure()