Package: theseus-ai
A library for differentiable nonlinear optimizationPaper•Blog•Webpage•Tutorials•DocsTheseus is an efficient application-agnostic library for building custom nonlinear optimization layers in PyTorch to support constructing various problems in robotics and vision as end-to-end differentiable architectures.Differentiable nonlinear optimization provides a general scheme to encode inductive priors, as the objective function can be partly parameterized by neural models and partly with expert domain-specific differentiable models. The ability to compute gradients end-to-end is retained by differentiating through the optimizer which allows neural models to train on the final task loss, while also taking advantage of priors captured by the optimizer.Current FeaturesApplication agnostic interfaceOur implementation provides an easy to use interface to build custom optimization layers and plug them into any neural architecture. Following differentiable features are currently available:Second-order nonlinear optimizersGauss-Newton (GN), Levenberg–Marquardt (LM), Trust Region, DoglegOther nonlinear optimizersCross Entropy Method (CEM)Linear solversDense: Cholesky, LU; Sparse: CHOLMOD, LU (GPU-only),BaSpaChoCommonly used costs,AutoDiffCostFunction,RobustCostFunctionLie groupsbased ontorchlieRobot kinematicsbased ontorchkinEfficiency based designWe support several features that improve computation times and memory consumption:Sparse linear solversBatching and GPU accelerationAutomatic vectorizationBackward modesImplicit, Truncated, Direct Loss Minimization (DLM), Sampling (LEO)Getting StartedPrerequisitesWestronglyrecommend you install Theseus in a venv or conda environment with Python 3.8-3.10.Theseus requirestorchinstallation. To install for your particular CPU/CUDA configuration, follow the instructions in the PyTorchwebsite.For GPU support, Theseus requiresnvccto compile custom CUDA operations. Make sure it matches the version used to compile pytorch withnvcc --version. If not, install it and ensure its location is on your system's$PATHvariable.Theseus also requiressuitesparse, which you can install via:sudo apt-get install libsuitesparse-dev(Ubuntu).conda install -c conda-forge suitesparse(Mac).Installingpypipipinstalltheseus-aiWe currently provide wheels with our CUDA extensions compiled using CUDA 11.6 and Python 3.10.
For other CUDA versions, consider installing from source or using ourbuild script.Note thatpypiinstallation doesn't include our experimentalTheseus Labs.
For this, please install from source.From sourceThe simplest way to install Theseus from source is by running the following (see further below to also include BaSpaCho)gitclonehttps://github.com/facebookresearch/theseus.git&&cdtheseus
pipinstall-e.If you are interested in contributing to Theseus, instead installpipinstall-e".[dev]"pre-commitinstalland follow the more detailed instructions inCONTRIBUTING.Installing BaSpaCho extensions from sourceBy default, installing from source doesn't include our BaSpaCho sparse solver extension. For this, follow these steps:Compile BaSpaCho from source following instructionshere. We recommend using flags-DBLA_STATIC=ON -DBUILD_SHARED_LIBS=OFF.Rungitclonehttps://github.com/facebookresearch/theseus.git&&cdtheseusBASPACHO_ROOT_DIR=<path/to/root/baspacho/dir>pipinstall-e.where the BaSpaCho root dir must have the binaries in the subdirectorybuild.Running unit tests (requiresdevinstallation)python-mpytesttestsBy default, unit tests include tests for our CUDA extensions. You can add the option-m "not cudaext"to skip them when installing without CUDA support. Additionally, the tests for sparse solver BaSpaCho are automatically
skipped when its extlib is not compiled.ExamplesSimple example. This example is fitting the curve $y$ to a dataset of $N$ observations $(x,y) \sim D$. This is modeled as anObjectivewith a singleCostFunctionthat computes the residual $y - v e^x$. TheObjectiveand theGaussNewtonoptimizer are encapsulated into aTheseusLayer. WithAdamand MSE loss, $x$ is learned by differentiating through theTheseusLayer.importtorchimporttheseusasthx_true,y_true,v_true=read_data()# shapes (1, N), (1, N), (1, 1)x=th.Variable(torch.randn_like(x_true),name="x")y=th.Variable(y_true,name="y")v=th.Vector(1,name="v")# a manifold subclass of Variable for optim_varsdeferror_fn(optim_vars,aux_vars):# returns y - v * exp(x)x,y=aux_varsreturny.tensor-optim_vars[0].tensor*torch.exp(x.tensor)objective=th.Objective()cost_function=th.AutoDiffCostFunction([v],error_fn,y_true.shape[1],aux_vars=[x,y],cost_weight=th.ScaleCostWeight(1.0))objective.add(cost_function)layer=th.TheseusLayer(th.GaussNewton(objective,max_iterations=10))phi=torch.nn.Parameter(x_true+0.1*torch.ones_like(x_true))outer_optimizer=torch.optim.Adam([phi],lr=0.001)forepochinrange(10):solution,info=layer.forward(input_tensors={"x":phi.clone(),"v":torch.ones(1,1)},optimizer_kwargs={"backward_mode":"implicit"})outer_loss=torch.nn.functional.mse_loss(solution["v"],v_true)outer_loss.backward()outer_optimizer.step()Seetutorials, and robotics and visionexamplesto learn about the API and usage.Citing TheseusIf you use Theseus in your work, please cite thepaperwith the BibTeX below.@article{pineda2022theseus,title={{Theseus: A Library for Differentiable Nonlinear Optimization}},author={Luis Pineda and Taosha Fan and Maurizio Monge and Shobha Venkataraman and Paloma Sodhi and Ricky TQ Chen and Joseph Ortiz and Daniel DeTone and Austin Wang and Stuart Anderson and Jing Dong and Brandon Amos and Mustafa Mukadam},journal={Advances in Neural Information Processing Systems},year={2022}}LicenseTheseus is MIT licensed. See theLICENSEfor details.Additional InformationJoin the community onGithub Discussionsfor questions and sugesstions.UseGithub Issuesfor bugs and features.SeeCONTRIBUTINGif interested in helping out.Theseus is made possible by the following contributors:Made withcontrib.rocks.