一个用于可微分非线性优化的库

论文 • 视频 • Twitter • 网页 • 教程

Theseus 是一个高效且与应用无关的库，用于在 PyTorch 中构建自定义的非线性优化层，以支持将机器人和视觉领域的各种问题构建成端到端可微分的架构。

可微分的非线性优化提供了一种通用的方法来编码归纳先验，因为目标函数可以部分由神经网络模型参数化，而另一部分则使用领域专家设计的可微分模型。通过在优化器中进行反向传播，仍然可以保留端到端计算梯度的能力，从而使神经网络能够在最终的任务损失上进行训练，同时还能利用优化器所捕捉到的先验信息。

请参阅 论文列表，其中包含了使用 Theseus 发表的、涵盖多个应用领域的示例。

---

当前特性

应用无关的接口

我们的实现提供了一个易于使用的接口，用于构建自定义的优化层，并将其插入到任何神经网络架构中。目前支持以下可微分功能：

• 二阶非线性优化器
  • 高斯-牛顿法 (GN)、莱文伯格-马夸尔特法 (LM)、信赖域法、狗腿法
• 其他非线性优化器
  • 交叉熵方法 (CEM)
• 线性求解器
  • 密集型：Cholesky 分解、LU 分解；稀疏型：CHOLMOD、LU（仅 GPU）、BaSpaCho
• 常用代价函数、自动微分代价函数、鲁棒代价函数
• 基于 torchlie 的 李群
• 基于 torchkin 的 机器人运动学

高效的设计

我们支持多项能够提升计算速度和降低内存消耗的功能：

• 稀疏线性求解器
• 批处理和 GPU 加速
• 自动向量化
• 反向模式
  • 隐式模式、截断模式、直接损失最小化 (DLM)、采样模式 (LEO)

入门指南

先决条件

• 我们强烈建议您在 Python 3.8–3.10 的 venv 或 conda 环境中安装 Theseus。
• Theseus 需要安装 torch。请根据您的 CPU/CUDA 配置，按照 PyTorch 官网 的说明进行安装。
• 为了支持 GPU，Theseus 需要 nvcc 来编译自定义的 CUDA 操作。请确保其版本与编译 PyTorch 时使用的 nvcc --version 版本一致。如果不一致，请安装并确保其路径已添加到系统的 $PATH 变量中。
• Theseus 还需要 suitesparse，您可以按以下方式安装：
  • sudo apt-get install libsuitesparse-dev（Ubuntu）。
  • conda install -c conda-forge suitesparse（Mac）。

安装

• pypi

  pip install theseus-ai
  

  目前我们提供的 wheel 包是使用 CUDA 11.6 和 Python 3.10 编译的 CUDA 扩展。对于其他 CUDA 版本，请考虑从源码安装，或使用我们的 构建脚本。

  请注意，通过 pypi 安装不会包含我们的实验性 Theseus Labs。 如果需要这些功能，请从源码安装。

• 从源码安装

  从源码安装 Theseus 的最简单方法是运行以下命令（请参阅下文以同时包含 BaSpaCho）：

  git clone https://github.com/facebookresearch/theseus.git && cd theseus
  pip install -e .
  

  如果您有兴趣为 Theseus 做贡献，可以安装：

  pip install -e ".[dev]"
  pre-commit install
  

  并按照 CONTRIBUTING 中更详细的说明进行操作。

• 从源码安装 BaSpaCho 扩展

  默认情况下，从源码安装不会包含我们的 BaSpaCho 稀疏求解器扩展。要包含该扩展，请按照以下步骤操作：

  1. 按照 这里 的说明从源码编译 BaSpaCho。我们建议使用 -DBLA_STATIC=ON -DBUILD_SHARED_LIBS=OFF 标志。

  2. 运行

     git clone https://github.com/facebookresearch/theseus.git && cd theseus
     BASPACHO_ROOT_DIR=<ba-spacho根目录路径> pip install -e .
     

     其中，BaSpaCho 的根目录必须在 build 子目录中包含二进制文件。

运行单元测试（需要 dev 安装）

python -m pytest tests

默认情况下，单元测试会包含对 CUDA 扩展的测试。如果未安装 CUDA 支持，可以通过添加 -m "not cudaext" 选项来跳过这些测试。此外，当 BaSpaCho 稀疏求解器扩展未编译时，其测试也会自动跳过。

示例

简单示例。此示例将曲线 $y$ 拟合到包含 $N$ 个观测值 $(x,y) \sim D$ 的数据集上。这被建模为一个 Objective，其中包含一个计算残差 $y - v e^x$ 的 CostFunction。Objective 和 GaussNewton 优化器被封装进一个 TheseusLayer 中。通过 Adam 优化器和 MSE 损失函数，$x$ 变量可以通过对 TheseusLayer 的反向传播来学习。

import torch
import theseus as th

x_true, y_true, v_true = read_data() # 形状分别为 (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") # 用于优化变量的流形子类 Variable

def error_fn(optim_vars, aux_vars): # 返回 y - v * exp(x)
    x, y = aux_vars
    return y.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)
for epoch in range(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()

更多关于 API 和用法的信息，请参阅 教程 以及机器人和视觉领域的 示例。

引用 Theseus

如果您在工作中使用了 Theseus，请引用以下论文，并使用提供的 BibTeX 格式：

@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}
}

许可证

Theseus 采用 MIT 许可证。详细信息请参阅 LICENSE。

更多信息

• 如有任何问题或建议，请加入 Github Discussions 社区。
• 如发现错误或有新功能需求，请使用 Github Issues 提交。
• 如果您有兴趣参与贡献，请参阅 CONTRIBUTING。

Theseus 的开发离不开以下贡献者的支持：

由 contrib.rocks 制作。

Theseus 快速上手指南

Theseus 是一个用于构建可微分非线性优化层的 PyTorch 库，专为机器人学和计算机视觉领域的端到端可微架构设计。它允许将专家先验知识与神经网络结合，通过可微优化器进行反向传播训练。

环境准备

在开始之前，请确保满足以下系统要求和依赖：

• Python 版本：推荐 3.8 - 3.10（强烈建议在 venv 或 conda 环境中安装）。
• PyTorch：必须预先安装。请根据您的 CPU/CUDA 配置参考 PyTorch 官网 进行安装。
  • 国内加速：可使用清华源安装，例如：pip install torch torchvision torchaudio --index-url https://pypi.tuna.tsinghua.edu.cn/simple
• GPU 支持（可选）：若需使用 GPU 加速自定义 CUDA 操作，需安装 nvcc。请确保 nvcc --version 显示的版本与编译当前 PyTorch 所用的 CUDA 版本一致，并将 nvcc 路径加入系统 $PATH。
• SuiteSparse：Theseus 依赖此库进行稀疏线性求解。
  • Ubuntu: sudo apt-get install libsuitesparse-dev
  • macOS (Conda): conda install -c conda-forge suitesparse
  • 国内加速：使用清华或中科大 Conda 源可显著提升下载速度。

安装步骤

方式一：通过 PyPI 安装（推荐）

这是最简单的安装方式，适用于大多数用户。预编译包基于 CUDA 11.6 和 Python 3.10。

pip install theseus-ai

注意：PyPI 版本不包含实验性的 "Theseus Labs" 模块。如需使用该部分功能，请选择源码安装。

方式二：从源码安装

如果您需要最新功能、贡献代码或使用实验性模块，请从源码安装。

1. 克隆仓库并进入目录：

   git clone https://github.com/facebookresearch/theseus.git && cd theseus
   

2. 安装开发依赖（包含测试和格式化工具）：

   pip install -e ".[dev]"
   pre-commit install
   

   若仅需基础功能，可运行：

   pip install -e .
   

（可选）安装 BaSpaCho 稀疏求解器扩展

默认源码安装不包含 BaSpaCho。如需启用该高性能稀疏求解器：

1. 按照 BaSpaCho 官方说明 从源码编译，推荐使用标志 -DBLA_STATIC=ON -DBUILD_SHARED_LIBS=OFF。
2. 设置环境变量并安装 Theseus：

   BASPACHO_ROOT_DIR=<path/to/root/baspacho/dir> pip install -e .
   

   (请将 <path/to/root/baspacho/dir> 替换为包含 build 子目录的 BaSpaCho 根路径)

基本使用

以下示例演示了如何使用 Theseus 拟合曲线 $y = v \cdot e^x$。我们将构建一个优化目标，将其封装为 TheseusLayer，并通过外部优化器（Adam）学习参数 $x$。

import torch
import theseus as th

# 假设已有数据读取函数
x_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") # 用于优化的流形变量子类

# 定义误差函数：residual = y - v * exp(x)
def error_fn(optim_vars, aux_vars):
    x, y = aux_vars
    return y.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)

# 封装为 TheseusLayer，使用 Gauss-Newton 优化器
layer = th.TheseusLayer(th.GaussNewton(objective, max_iterations=10))

# 初始化待学习参数 phi (对应 x)
phi = torch.nn.Parameter(x_true + 0.1 * torch.ones_like(x_true))
outer_optimizer = torch.optim.Adam([phi], lr=0.001)

# 训练循环
for epoch in range(10):
    # 前向传播：执行优化
    solution, info = layer.forward(
        input_tensors={"x": phi.clone(), "v": torch.ones(1, 1)},
        optimizer_kwargs={"backward_mode": "implicit"})
    
    # 计算外部损失 (MSE)
    outer_loss = torch.nn.functional.mse_loss(solution["v"], v_true)
    
    # 反向传播并更新参数
    outer_loss.backward()
    outer_optimizer.step()

更多详细教程和应用案例（如机器人运动学、状态估计等），请参阅官方 Tutorials 和 Examples 目录。