安装 | 文档 | 教程 | 示例 | 论文 | 引用

TorchOpt 是一个基于 PyTorch 构建的高效可微优化库。 TorchOpt 具有以下特点：

• 全面性：TorchOpt 提供了三种不同的微分模式——显式微分、隐式微分和零阶微分，以应对各种可微优化场景。
• 灵活性：TorchOpt 同时提供了函数式和面向目标的 API，满足用户的不同偏好。用户可以采用类似 JAX 或 PyTorch 的风格实现可微优化。
• 高效性：TorchOpt 提供了（1）CPU/GPU 加速的可微优化器；（2）基于 RPC 的分布式训练框架；（3）快速树运算，从而大幅提高双层优化问题的训练效率。

除了可微优化之外，TorchOpt 还可以被视为一种函数式优化器，支持类似 JAX 的可组合函数式优化器用于 PyTorch。 借助 TorchOpt，用户可以轻松地以函数式风格在 PyTorch 中进行神经网络优化，类似于 JAX 中的 Optax。

---

本 README 的组织结构如下：

• TorchOpt 作为函数式优化器
  • 类似 Optax 的 API
  • 类似 PyTorch 的 API
  • 可微性
• TorchOpt 用于可微优化
  • 显式梯度 (EG)
  • 隐式梯度 (IG)
  • 零阶微分 (ZD)
• 高性能与分布式训练
  • CPU/GPU 加速的可微优化器
  • 分布式训练
  • OpTree
• 可视化
• 示例
• 安装
• 更新日志
• 引用 TorchOpt
• 团队
• 许可

---

TorchOpt 作为函数式优化器

TorchOpt 的设计遵循函数式编程的理念。 与 functorch 保持一致，用户可以在 PyTorch 中使用模型、优化器和训练进行函数式编程。 我们以 Adam 优化器为例进行说明。 您还可以查看教程笔记本 函数式优化器 以获取更多细节。

类似Optax的API

对于偏好全功能编程的用户，我们提供类似Optax的API，只需将梯度和优化器状态传递给优化器函数即可。 以下是与functorch结合使用的示例：

class Net(nn.Module): ...

class Loader(DataLoader): ...

net = Net()  # 初始化
loader = Loader()
optimizer = torchopt.adam()

model, params = functorch.make_functional(net)           # 使用functorch提取网络参数
opt_state = optimizer.init(params)                       # 初始化优化器

xs, ys = next(loader)                                    # 获取数据
pred = model(params, xs)                                 # 前向传播
loss = F.cross_entropy(pred, ys)                         # 计算损失

grads = torch.autograd.grad(loss, params)                # 计算梯度
updates, opt_state = optimizer.update(grads, opt_state)  # 获取更新
params = torchopt.apply_updates(params, updates)         # 更新网络参数

我们还提供了一个包装器torchopt.FuncOptimizer，以简化优化器状态的维护：

net = Net()  # 初始化
loader = Loader()
optimizer = torchopt.FuncOptimizer(torchopt.adam())      # 使用torchopt.FuncOptimizer包装

model, params = functorch.make_functional(net)           # 使用functorch提取网络参数

for xs, ys in loader:                                    # 获取数据
    pred = model(params, xs)                             # 前向传播
    loss = F.cross_entropy(pred, ys)                     # 计算损失

    params = optimizer.step(loss, params)                # 更新网络参数

类似PyTorch的API

我们还设计了一个基类torchopt.Optimizer，其接口与torch.optim.Optimizer完全相同。 我们通过包装我们的类似Optax的API，为传统PyTorch用户提供原生PyTorch API（例如zero_grad()或step()）。

net = Net()  # 初始化
loader = Loader()
optimizer = torchopt.Adam(net.parameters())

xs, ys = next(loader)             # 获取数据
pred = net(xs)                    # 前向传播
loss = F.cross_entropy(pred, ys)  # 计算损失

optimizer.zero_grad()             # 清零梯度
loss.backward()                   # 反向传播
optimizer.step()                  # 执行更新

可微分性

除了与torch.optim相同的优化函数外，函数式优化器的一个重要优势是能够轻松实现可微分优化。 这一点在算法需要通过优化更新进行微分时尤其有用（例如元学习实践）。 我们以梯度和优化器状态作为输入，并使用非就地操作符来计算并输出更新。 这些过程可以自动实现，用户只需在函数中传入参数inplace=False即可。 请查看显式梯度（EG）函数式API部分以获取示例。

---

TorchOpt用于可微分优化

我们设计了一种双层优化更新方案，该方案可轻松扩展以实现各种可微分优化过程。

如上图所示，该方案包含一个外层，其参数$\phi$可通过内层参数解$\theta^{\prime}(\phi)$端到端地学习，利用最优响应导数$\partial \theta^{\prime}(\phi) / \partial \phi$。 TorchOpt支持三种微分模式。 可以看出，该算法的关键组件是计算最优响应（BR）雅可比矩阵。 从基于最优响应的角度来看，现有的梯度方法可分为三类：对展开优化路径的显式梯度、隐式微分以及零阶梯度微分。

显式梯度（EG）

显式梯度的思想是将梯度步视为一个可微分函数，并尝试通过展开的优化路径进行反向传播。 这种微分模式适用于内层优化解由少量梯度步得到的算法，例如MAML和MGRL。 TorchOpt同时提供了函数式和面向对象的API，以适应不同的用户应用。

函数式API

函数式API以函数式编程风格进行优化。 请注意，我们向函数传递了参数inplace=False，以使优化过程可微分。 更多指导请参阅教程笔记本函数式优化器。

# 定义函数式优化器
optimizer = torchopt.adam()
# 定义元参数和内层参数
meta_params = ...
fmodel, params = make_functional(model)
# 初始状态
state = optimizer.init(params)

for iter in range(iter_times):
    loss = inner_loss(fmodel, params, meta_params)
    grads = torch.autograd.grad(loss, params)
    # 应用非就地参数更新
    updates, state = optimizer.update(grads, state, inplace=False)
    params = torchopt.apply_updates(params, updates)

loss = outer_loss(fmodel, params, meta_params)
meta_grads = torch.autograd.grad(loss, meta_params)

面向对象API

TorchOpt还提供了兼容PyTorch编程风格的面向对象API。 更多指导请参阅示例和教程笔记本元优化器、停止梯度。

# 定义元参数和内层参数
meta_params = ...
model = ...
# 定义可微分优化器
optimizer = torchopt.MetaAdam(model)  # 将模型实例作为参数，而非model.parameters()

for iter in range(iter_times):
    # 执行内层更新
    loss = inner_loss(model, meta_params)
    optimizer.step(loss)

loss = outer_loss(model, meta_params)
loss.backward()

隐式梯度（IG）

通过将解 $\theta^{\prime}$ 视为 $\phi$ 的隐函数，隐式梯度的思想是直接利用隐函数定理求得解析的最优响应导数 $\partial \theta^{\prime} (\phi) / \partial \phi$。

这种方法适用于内层优化解达到以下条件的情况：$\left. \frac{\partial F (\theta, \phi)}{\partial \theta} \right\rvert_{\theta=\theta^{\prime}} = 0$ 或者满足某些驻点条件 $F (\theta^{\prime}, \phi) = 0$，例如 iMAML 和 DEQ。

TorchOpt 提供了函数式和面向对象两种 API，以支持基于共轭梯度法和Neumann 级数法的隐式梯度方法。

更多指导请参阅示例 iMAML 和笔记本 隐式梯度。

函数式 API

对于隐式梯度，与 JAXopt 类似，用户需要定义驻点条件，TorchOpt 提供了一个装饰器来包装求解函数，从而启用隐式梯度计算。

# 内层循环的驻点条件
def stationary(params, meta_params, data):
    # 驻点条件构建
    return 驻点条件

# 包装函数的装饰器
# 可选指定线性求解器（共轭梯度或 Neumann 级数）
@torchopt.diff.implicit.custom_root(stationary, solve=linear_solver)
def solve(params, meta_params, data):
    # 参数的前向优化过程
    return 输出

# 定义参数、元参数并获取数据
params, meta_prams, data = ..., ..., ...
optimal_params = solve(params, meta_params, data)
loss = outer_loss(optimal_params)

meta_grads = torch.autograd.grad(loss, meta_params)

面向对象 API

TorchOpt 还提供了面向对象的 API，用户需要继承自类 torchopt.nn.ImplicitMetaGradientModule 来构造内层网络。

用户需要定义驻点条件/目标函数以及内层循环的求解函数，以启用隐式梯度计算。

# 继承自类 ImplicitMetaGradientModule
# 可选指定线性求解器（共轭梯度或 Neumann 级数）
class InnerNet(ImplicitMetaGradientModule, linear_solve=linear_solver):
    def __init__(self, meta_param):
        super().__init__()
        self.meta_param = meta_param
        ...

    def forward(self, batch):
        # 前向过程
        ...

    def optimality(self, batch, labels):
        # 构建驻点条件以计算隐式梯度
        # 注意：如果未实现此方法，则会自动从 `objective` 函数的梯度推导出来。
        ...

    def objective(self, batch, labels):
        # 定义内层循环的优化目标
        ...

    def solve(self, batch, labels):
        # 执行内层循环的优化
        ...

# 获取元参数和数据
meta_params, data = ..., ...
inner_net = InnerNet(meta_params)

# 求解与元参数相关的内层循环过程
optimal_inner_net = inner_net.solve(data)

# 获取外层损失并求解元梯度
loss = outer_loss(optimal_inner_net)
meta_grads = torch.autograd.grad(loss, meta_params)

零阶微分（ZD）

当内层循环过程不可微或者希望消除前两种模式中由 Hessian 带来的繁重计算负担时，可以选择零阶微分（ZD）。

零阶微分通常基于零阶估计获取梯度，例如有限差分或 进化策略。

与优化目标 $F$ 不同，进化策略优化的是一个平滑后的目标。

TorchOpt 提供了函数式和面向对象两种 API，用于进化策略方法。

更多指导请参阅教程笔记本 零阶微分。

函数式 API

对于零阶微分，用户需要定义前向计算过程和噪声采样程序。TorchOpt 提供了一个装饰器来包装前向函数，从而启用零阶微分。

# 自定义进化策略中的噪声采样函数
def distribution(sample_shape):
    # 生成一批噪声样本
    # 注意：分布应为球对称且方差恒为 1。
    ...
    return 噪声批次

# 分布也可以是 `torch.distributions.Distribution` 的实例，例如 `torch.distributions.Normal(...)`
distribution = torch.distributions.Normal(loc=0, scale=1)

# 指定进化策略的方法和超参数
@torchopt.diff.zero_order(distribution, method)
def forward(params, batch, labels):
    # 前向过程
    ...
    return 目标值  # 返回的张量应为标量张量

面向对象 API

TorchOpt 还提供了面向对象的 API，用户需要继承自类 torchopt.nn.ZeroOrderGradientModule，以按照经典的 PyTorch 风格将网络构造为 nn.Module。

用户需要定义零阶梯度的前向过程 forward() 和噪声采样函数 sample()。

# 继承自类 ZeroOrderGradientModule
# 可选指定使用的 `method` 和/或 `num_samples` 和/或 `sigma` 用于采样
class Net(ZeroOrderGradientModule, method=method, num_samples=num_samples, sigma=sigma):
    def __init__(self, ...):
        ...

    def forward(self, batch):
        # 前向过程
        ...
        return 目标值  # 返回的张量应为标量张量

    def sample(self, sample_shape=torch.Size()):
        # 生成一批噪声样本
        # 注意：分布应为球对称且方差恒为 1。
        ...
        return 噪声批次

# 获取模型和数据
net = Net(...)
data = ...

# 前向传播
loss = Net(data)
# 使用零阶微分进行反向传播
grads = torch.autograd.grad(loss, net.parameters())

---

高性能与分布式训练

CPU/GPU加速的可微优化器

我们把优化器作为一个整体来处理，而不是将其拆分为若干基本算子（例如 sqrt 和 div）。 因此，通过手动编写前向和反向函数，我们可以进行符号化简化。 此外，我们还可以存储一些中间数据，这些数据在反向传播过程中可以被重复使用。 我们用 C++ OpenMP 和 CUDA 编写了加速函数，并通过 pybind11 将它们绑定，以便能够从 Python 中调用，然后使用 torch.autograd.Function 定义前向和反向行为。 用户只需将 use_accelerated_op 标志设为 True 即可使用该优化器。 请参考教程中的相应章节 功能优化器 和 元优化器。

optimizer = torchopt.MetaAdam(model, lr, use_accelerated_op=True)

分布式训练

TorchOpt 提供了基于 PyTorch RPC 模块的分布式训练功能，以实现更快的训练速度并支持多节点多 GPU。 与类似 MPI 的并行化范式不同，后者使用多个同质工作进程并需要精心设计的通信钩子，RPC API 允许用户更灵活地构建自己的优化流水线。 实验结果表明，我们实现了加速比与工作进程数量之间的近似线性关系。 请查看 分布式训练文档 和 分布式 MAML 示例，获取更具体的指导。

OpTree

我们实现了 PyTree，以利用 C++ 实现快速的嵌套结构扁平化。 树操作（例如扁平化和反扁平化）对于实现深度学习框架的功能性和即时编译（JIT）特性至关重要。 通过在 C++ 中实现，我们可以使用一些缓存和内存友好的数据结构（例如 absl::InlinedVector），从而提升性能。 如需更多指导和对比结果，请参阅我们的开源项目 OpTree。

---

可视化

元学习中复杂的梯度流给梯度流的管理和正确性验证带来了巨大挑战。 TorchOpt 提供了一个可视化工具，可在梯度图上标注变量（例如网络参数或元参数）名称，以便更好地分析。 该可视化工具基于 torchviz 修改而来。 请参阅示例 可视化代码 和教程笔记本 可视化，了解更多信息。

下图展示了可视化结果。 与 torchviz 相比，TorchOpt 将 Adam 内的操作融合在一起（橙色），以降低复杂度并提供更简单的可视化效果。

---

示例

在 examples 目录中，我们提供了几个使用 TorchOpt 的功能优化器示例和轻量级元学习示例。

• 模型无关元学习 (MAML) - 监督学习 (ICML 2017)
• 学习重新加权样本以实现鲁棒深度学习 (ICML 2018)
• 模型无关元学习 (MAML) - 强化学习 (ICML 2017)
• 元梯度强化学习 (MGRL) (NeurIPS 2018)
• 通过对手学习过程学习 (LOLA) (AAMAS 2018)
• 带有隐式梯度的元学习 (NeurIPS 2019)

此外，也请查看 examples 中更多分布式/可视化/与 functorch 兼容的示例。

---

安装

要求

• PyTorch
• （可选）用于可视化计算图
  • Graphviz（对于 Linux 用户，使用 apt/yum install graphviz 或 conda install -c anaconda python-graphviz）

请先按照 https://pytorch.org 上的说明，在您的 Python 环境中安装 PyTorch。 然后运行以下命令，从 PyPI 安装 TorchOpt（ / )：

pip3 install torchopt

如果 PyTorch 的最低版本不满足要求，pip 会为您安装或升级。请注意 torch 的构建版本，以支持 CPU 或 CUDA（例如 cpu、cu118、cu121）。 您可能需要指定额外的索引 URL 来安装 torch 包：

pip3 install torchopt --extra-index-url https://download.pytorch.org/whl/cu121

有关安装 PyTorch 的更多信息，请参阅 https://pytorch.org。

您也可以从源码构建共享库，使用：

git clone https://github.com/metaopt/torchopt.git
cd torchopt
pip3 install .

我们提供了一个 conda 环境配方，用于安装构建工具链，例如 cmake、g++ 和 nvcc。 您可以使用以下命令，借助 conda / mamba 创建一个新的隔离环境。

git clone https://github.com/metaopt/torchopt.git
cd torchopt

# 如果 conda 无法检测到 NVIDIA 驱动程序（例如在 docker 或 WSL2 中），您可能需要设置 `CONDA_OVERRIDE_CUDA`
CONDA_OVERRIDE_CUDA=12.1 conda env create --file conda-recipe-minimal.yaml

conda activate torchopt
make install-editable  # 或运行 `pip3 install --no-build-isolation --editable .`

---

更改日志

请参阅 CHANGELOG.md。

---

引用 TorchOpt

如果您觉得 TorchOpt 有用，请在您的论文中引用它。

@article{JMLR:TorchOpt,
  author  = {Jie Ren* and Xidong Feng* and Bo Liu* and Xuehai Pan* and Yao Fu and Luo Mai and Yaodong Yang},
  title   = {TorchOpt: 一个高效的可微优化库},
  journal = {机器学习研究杂志},
  year    = {2023},
  volume  = {24},
  number  = {367},
  pages   = {1--14},
  url     = {http://jmlr.org/papers/v24/23-0191.html}
}

团队

TorchOpt 由 Jie Ren、Xidong Feng、Bo Liu、Xuehai Pan、Luo Mai 和 Yaodong Yang 共同完成。

许可协议

TorchOpt 采用 Apache 许可协议第 2.0 版发布。

TorchOpt 快速上手指南

环境准备

• Python 版本：3.8+
• 前置依赖：PyTorch（推荐 2.0+）
• 推荐使用 conda 或 pip 管理环境，支持 CPU/GPU 加速

建议使用国内镜像源加速安装：
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torchopt

安装步骤

pip install torchopt

如需使用最新开发版，可从 GitHub 安装：

pip install git+https://github.com/metaopt/torchopt.git

基本使用

方式一：Functional 风格（推荐用于可微优化）

import torch
import torch.nn as nn
import torchopt
from functorch import make_functional

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(2, 1)

net = Net()
optimizer = torchopt.adam()  # 创建优化器

model, params = make_functional(net)  # 转为函数式模型
opt_state = optimizer.init(params)    # 初始化状态

# 模拟训练步骤
xs = torch.randn(4, 2)
ys = torch.randn(4, 1)
pred = model(params, xs)
loss = torch.nn.functional.mse_loss(pred, ys)

grads = torch.autograd.grad(loss, params)
updates, opt_state = optimizer.update(grads, opt_state, inplace=False)  # 非原地更新，支持自动微分
params = torchopt.apply_updates(params, updates)

方式二：PyTorch 风格（传统接口）

import torch
import torch.nn as nn
import torchopt

net = nn.Linear(2, 1)
optimizer = torchopt.Adam(net.parameters())  # 与 torch.optim.Adam 接口一致

xs = torch.randn(4, 2)
ys = torch.randn(4, 1)

pred = net(xs)
loss = torch.nn.functional.mse_loss(pred, ys)

optimizer.zero_grad()
loss.backward()
optimizer.step()

✅ 推荐初学者从 PyTorch 风格 开始，进阶用户使用 Functional 风格 实现元学习、可微优化等高级场景。