TorchOpt 优化器

Optimizer(params, impl)

一个类似于 torch.optim.Optimizer 的经典优化器的基类。

MetaOptimizer(module, impl)

高级可微分优化器的基类。

优化器

class torchopt.Optimizer(params, impl)[source]

基础类: object

一个类似于torch.optim.Optimizer的经典优化器的基类。

初始化优化器。

Parameters:
  • params (可迭代的 torch.Tensor) – 一个torch.Tensor的可迭代对象。指定应该优化的张量。

  • impl (GradientTransformation) – 一个低级别的优化器函数,它可以是torchopt.alias中提供的优化器函数,也可以是一个自定义的torchopt.chain()转换。 注意,使用Optimizer(sgd())Optimizer(chain(sgd()))等同于torchopt.SGD

__init__(params, impl)[source]

初始化优化器。

Parameters:
  • params (可迭代的 torch.Tensor) – 一个torch.Tensor的可迭代对象。指定应该优化的张量。

  • impl (GradientTransformation) – 一个低级别的优化器函数,它可以是torchopt.alias中提供的优化器函数,也可以是一个自定义的torchopt.chain()转换。 注意,使用Optimizer(sgd())Optimizer(chain(sgd()))等同于torchopt.SGD

zero_grad(set_to_none=False)[source]

将所有优化的torch.Tensor的梯度设置为零。

行为类似于 torch.optim.Optimizer.zero_grad()

Parameters:

set_to_none (bool, 可选) – 不是设置为零,而是将 grads 设置为 None。(默认值:False

Return type:

state_dict()[source]

返回优化器的状态。

Return type:

tuple[OptState, …]

load_state_dict(state_dict)[source]

加载优化器状态。

Parameters:

state_dict (sequence of tree of Tensor) – 优化器状态。应该是从调用state_dict()返回的对象。

Return type:

step(closure=None)[source]

执行单个优化步骤。

行为类似于 torch.optim.Optimizer.step()

Parameters:

closure (callableNone, 可选) – 一个重新评估模型并返回损失的闭包。对于大多数优化器来说是可选的。(默认值:None

Return type:

Tensor | None

add_param_group(params)[source]

向优化器的param_groups添加一个参数组。

Return type:

MetaOptimizer

class torchopt.MetaOptimizer(module, impl)[source]

基础类: object

高级可微分优化器的基类。

初始化元优化器。

Parameters:
  • module (nn.Module) – 一个需要优化参数的网络。

  • impl (GradientTransformation) – 一个低级别的优化器函数,它可以是torchopt.alias中提供的优化器函数,也可以是一个自定义的torchopt.chain()转换。 注意,使用MetaOptimizer(sgd(moment_requires_grad=True))MetaOptimizer(chain(sgd(moment_requires_grad=True)))等同于 torchopt.MetaSGD

__init__(module, impl)[source]

初始化元优化器。

Parameters:
  • module (nn.Module) – 一个需要优化参数的网络。

  • impl (GradientTransformation) – 一个低级别的优化器函数,它可以是torchopt.alias中提供的优化器函数,也可以是一个自定义的torchopt.chain()转换。 注意,使用MetaOptimizer(sgd(moment_requires_grad=True))MetaOptimizer(chain(sgd(moment_requires_grad=True)))等同于 torchopt.MetaSGD

step(loss)[source]

计算损失对网络参数的梯度并更新网络参数。

将构建导数的图形,允许计算更高阶的导数乘积。我们使用可微分的优化器(传递参数inplace=False)来缩放梯度并更新网络参数,而不就地修改张量。

Parameters:

loss (torch.Tensor) – 用于计算网络参数梯度的损失。

Return type:

add_param_group(module)[source]

向优化器的state_groups添加一个参数组。

Return type:

state_dict()[source]

提取优化器状态的引用。

请注意,状态是引用,因此任何就地操作都会同时更改MetaOptimizer内部的状态。

Return type:

tuple[OptState, …]

load_state_dict(state_dict)[source]

加载优化器状态的引用。

Return type:


功能优化器

FuncOptimizer(impl, *[, inplace])

一个包装类,用于持有功能优化器。

adadelta([lr, rho, eps, weight_decay, ...])

创建一个AdaDelta优化器的功能版本。

adagrad([lr, lr_decay, weight_decay, ...])

创建一个AdaGrad优化器的功能版本。

adam([lr, betas, eps, weight_decay, ...])

创建一个功能版本的Adam优化器。

adamw([lr, betas, eps, weight_decay, ...])

创建一个带有权重衰减正则化的Adam优化器的功能版本。

adamax([lr, betas, eps, weight_decay, ...])

创建一个AdaMax优化器的功能版本。

radam([lr, betas, eps, weight_decay, ...])

创建一个功能版本的RAdam优化器。

rmsprop([lr, alpha, eps, weight_decay, ...])

创建一个RMSProp优化器的功能版本。

sgd(lr[, momentum, dampening, weight_decay, ...])

创建一个标准随机梯度下降优化器的功能版本。

函数优化器的包装器

class torchopt.FuncOptimizer(impl, *, inplace=False)[source]

基础类: object

一个包装类,用于保存功能优化器。

这个包装器使得维护优化器状态变得更加容易。优化器状态由包装器内部持有。包装器提供了一个step()函数来计算梯度并更新参数。

另请参阅

初始化功能优化器包装器。

Parameters:
  • impl (GradientTransformation) – 一个低级别的优化器函数,它可以是torchopt.alias中提供的优化器函数,也可以是一个自定义的torchopt.chain()转换。

  • inplace (bool, 可选) – 每次优化更新的inplace的默认值。 (默认: False)

__init__(impl, *, inplace=False)[source]

初始化功能优化器包装器。

Parameters:
  • impl (GradientTransformation) – 一个低级别的优化器函数,它可以是torchopt.alias中提供的优化器函数,也可以是一个自定义的torchopt.chain()转换。

  • inplace (bool, 可选) – 每次优化更新的inplace默认值。 (默认: False)

step(loss, params, inplace=None)[source]

计算损失对网络参数的梯度并更新网络参数。

将构建导数的图形,允许计算更高阶的导数乘积。我们使用可微分的优化器(传递参数inplace=False)来缩放梯度并更新网络参数,而不就地修改张量。

Parameters:
  • loss (Tensor) – 用于计算网络参数梯度的损失。

  • params (tree of Tensor) – 一个由 torch.Tensor 组成的树。指定了应该优化的张量。

  • inplace (boolNone, 可选) – 是否就地更新参数。如果 None,则使用构造函数中指定的默认值。 (默认: None)

Return type:

参数

state_dict()[source]

提取优化器状态的引用。

请注意,状态是引用,因此任何原地操作都会同时改变FuncOptimizer内部的状态。

Return type:

OptState

load_state_dict(state_dict)[source]

加载优化器状态的引用。

Return type:

功能性AdaDelta优化器

torchopt.adadelta(lr=0.001, rho=0.9, eps=1e-06, weight_decay=0.0, *, moment_requires_grad=False)[source]

创建一个AdaDelta优化器的功能版本。

Adadelta 是一种用于梯度下降的逐维度学习率方法。

参考文献

Parameters:
  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认: 1e-3)

  • rho (float, 可选) – 用于计算梯度及其平方的运行平均值的系数。 (默认值: 0.9)

  • eps (float, 可选) – 一个小的常数,应用于平方根(如Adadelta论文中所述),以避免在重新缩放时除以零。 (默认: 1e-6)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建为 标志 requires_grad=True,此标志常用于元学习算法。 (默认: False)

Return type:

GradientTransformation

Returns:

对应的 GradientTransformation 实例。

另请参阅

功能优化器包装器 torchopt.FuncOptimizer

功能性AdaGrad优化器

torchopt.adagrad(lr=0.01, lr_decay=0.0, weight_decay=0.0, initial_accumulator_value=0.0, eps=1e-10, *, maximize=False)[source]

创建一个AdaGrad优化器的功能版本。

AdaGrad 是一种基于梯度的优化算法,它在训练过程中为每个参数调整学习率。

警告

AdaGrad的主要限制是分母中平方梯度的单调累积。 由于所有项都是> 0,训练过程中总和不断增长,学习率最终变得非常小。

参考文献

Duchi 等人, 2011: https://jmlr.org/papers/v12/duchi11a.html

Parameters:
  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认值: 1e-2)

  • lr_decay (float, 可选) – 学习率衰减。(默认: 0.0)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • initial_accumulator_value (float, 可选) – 累加器的初始值。 (默认: 0.0)

  • eps (float, 可选) – 一个小的常数,应用于平方根外的分母(如Adam论文中所述),以避免在重新缩放时除以零。 (默认: 1e-10)

  • 最大化 (布尔值, 可选) – 根据目标最大化参数,而不是最小化。 (默认值: False)

Return type:

GradientTransformation

Returns:

对应的 GradientTransformation 实例。

另请参阅

功能优化器包装器 torchopt.FuncOptimizer

功能性Adam优化器

torchopt.adam(lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, *, eps_root=0.0, moment_requires_grad=False, maximize=False, use_accelerated_op=False)[source]

创建一个Adam优化器的功能版本。

Adam 是一种具有学习率自适应功能的 SGD 变体。用于每个权重的学习率是根据梯度的一阶和二阶矩的估计值(使用适当的指数移动平均)计算得出的。

参考文献

Parameters:
  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认: 1e-3)

  • betas (tuple of float, optional) – 用于计算梯度及其平方的运行平均值的系数。(默认值: (0.9, 0.999))

  • eps (float, 可选) – 一个小的常数,应用于平方根外的分母(如Adam论文中所述),以避免在重新缩放时除以零。 (默认: 1e-8)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • eps_root (float, optional) – 在平方根内的分母上应用的一个小常数(如RMSProp中所示),以避免在重新缩放时除以零。例如,在通过Adam计算(元)梯度时需要这个常数。(默认值:0.0

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建为 标志 requires_grad=True,此标志常用于元学习算法。 (默认: False)

  • 最大化 (布尔值, 可选) – 根据目标最大化参数,而不是最小化。 (默认: False)

  • use_accelerated_op (bool, 可选) – 如果 True 使用我们实现的融合操作符。 (默认: False)

Return type:

GradientTransformation

Returns:

对应的 GradientTransformation 实例。

另请参阅

功能优化器包装器 torchopt.FuncOptimizer

功能性AdamW优化器

torchopt.adamw(lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.01, *, eps_root=0.0, mask=None, moment_requires_grad=False, maximize=False, use_accelerated_op=False)[source]

创建一个带有权重衰减正则化的Adam优化器的功能版本。

AdamW 使用权重衰减来规范学习过程,使其倾向于较小的权重,因为这有助于更好的泛化。在 SGD 中,你也可以使用 L2 正则化来实现这一点,作为一个附加的损失项,然而 L2 正则化对于像 Adam 这样的自适应梯度算法并不如预期那样有效。

参考文献

Parameters:
  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认: 1e-3)

  • betas (tuple of float, optional) – 用于计算梯度及其平方的运行平均值的系数。(默认值: (0.9, 0.999))

  • eps (float, optional) – 一个小的常数,应用于平方根外的分母(如Adam论文中所述),以避免在重新缩放时除以零。 (默认值: 1e-8)

  • weight_decay (float, 可选) – 权重衰减正则化的强度。请注意,这个权重衰减会与学习率相乘。这与PyTorch等其他框架一致,但与(Loshchilov et al., 2019)不同,后者中权重衰减仅与“计划乘数”相乘,而不与基础学习率相乘。(默认: 1e-2)

  • eps_root (float, optional) – 一个小的常数,应用于平方根内的分母(如在RMSProp中),以避免在重新缩放时除以零。这在通过Adam计算(元)梯度时是必需的。(默认值:0.0

  • mask (tree of Tensor, callable, or None, optional) – 一个与参数 pytree 结构相同(或为其前缀)的树,或者一个根据参数/更新返回此类 pytree 的函数。叶子节点应为布尔值, True 表示您希望应用权重衰减的叶子/子树,而 False 表示您希望跳过的叶子/子树。请注意,Adam 梯度变换应用于所有参数。(默认值:None

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建为 标志 requires_grad=True,此标志常用于元学习算法。 (默认: False)

  • maximize (bool, optional) – 根据目标最大化参数,而不是最小化。(默认: False)

  • use_accelerated_op (bool, 可选) – 如果 True 使用我们实现的融合操作符。 (默认: False)

Return type:

GradientTransformation

Returns:

对应的 GradientTransformation 实例。

另请参阅

功能优化器包装器 torchopt.FuncOptimizer

功能性AdaMax优化器

torchopt.adamax(lr=0.001, betas=(0.9, 0.999), eps=1e-06, weight_decay=0.0, *, moment_requires_grad=False)[source]

创建一个AdaMax优化器的功能版本。

参考文献

Parameters:
  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认: 1e-3)

  • betas (tuple of float, optional) – 用于计算梯度及其平方的运行平均值的系数。(默认值: (0.9, 0.999))

  • eps (float, 可选) – 一个小的常数,应用于平方根(如RAdam论文中所述),以避免在重新缩放时除以零。 (默认: 1e-6)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建为 标志 requires_grad=True,此标志常用于元学习算法。 (默认: False)

Return type:

GradientTransformation

Returns:

对应的 GradientTransformation 实例。

另请参阅

功能优化器包装器 torchopt.FuncOptimizer

功能性RAdam优化器

torchopt.radam(lr=0.001, betas=(0.9, 0.999), eps=1e-06, weight_decay=0.0, *, moment_requires_grad=False)[source]

创建一个RAdam优化器的功能版本。

RAdam 是一种自适应学习率修正优化器的变体。

参考文献

Parameters:
  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认: 1e-3)

  • betas (tuple of float, optional) – 用于计算梯度及其平方的运行平均值的系数。(默认值: (0.9, 0.999))

  • eps (float, 可选) – 一个应用于平方根的小常数(如RAdam论文中所述),以避免在重新缩放时除以零。 (默认: 1e-6)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建为 标志 requires_grad=True,此标志常用于元学习算法。 (默认: False)

Return type:

GradientTransformation

Returns:

对应的 GradientTransformation 实例。

另请参阅

功能优化器包装器 torchopt.FuncOptimizer

功能性RMSProp优化器

torchopt.rmsprop(lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0.0, momentum=0.0, centered=False, *, initial_scale=0.0, nesterov=False, maximize=False)[source]

创建一个RMSProp优化器的功能版本。

RMSProp 是一种具有学习率自适应功能的 SGD 变体。每个权重使用的学习率会根据先前步骤中梯度大小的适当估计进行缩放。文献中可以找到几种 RMSProp 的变体。此别名提供了一个易于配置的 RMSProp 优化器,可用于在这些变体之间切换。

参考文献

Parameters:
  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认值: 1e-2)

  • alpha (float, 可选) – 平滑常数,用于跟踪先前梯度幅度的衰减。(默认: 0.99)

  • eps (float, 可选) – 一个小的数值常数,用于在重新缩放时避免除以零。 (默认: 1e-8)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • 动量 (浮点数, 可选) – 动量项使用的衰减率。当设置为 0.0 时不使用动量。(默认值: 0.0)

  • 居中 (bool, 可选) – 如果 True,使用过去梯度的方差来重新缩放最新的梯度。(默认: False)

  • initial_scale (float, optional) – 用于初始化跟踪先前更新幅度的累加器。PyTorch 使用 0.0,TensorFlow 1.x 使用 1.0。在复现论文结果时,请验证作者使用的值。 (默认值: 0.0)

  • nesterov (bool, 可选) – 是否使用Nesterov动量。(默认: False)

  • maximize (bool, optional) – 根据目标最大化参数,而不是最小化。(默认: False)

Return type:

GradientTransformation

Returns:

对应的 GradientTransformation 实例。

另请参阅

功能优化器包装器 torchopt.FuncOptimizer

功能性SGD优化器

torchopt.sgd(lr, momentum=0.0, dampening=0.0, weight_decay=0.0, nesterov=False, *, moment_requires_grad=False, maximize=False)[source]

创建一个规范随机梯度下降优化器的功能版本。

这实现了随机梯度下降。它还包括对动量和Nesterov加速的支持,因为这些是使用随机梯度下降训练深度神经网络时的标准做法。

参考文献

Parameters:
  • lr (floatcallable) – 这是一个固定的全局缩放因子或学习率调度器。

  • 动量 (float, 可选) – 动量项使用的衰减率。当设置为0.0时不使用动量。(默认: 0.0)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • dampening (float, 可选) – 动量的阻尼系数。(默认: 0.0)

  • nesterov (bool, 可选) – 是否使用Nesterov动量。(默认: False)

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建为 标志 requires_grad=True,此标志常用于元学习算法。 (默认: False)

  • maximize (bool, optional) – 根据目标最大化参数,而不是最小化。(默认: False)

Return type:

GradientTransformation

Returns:

对应的 GradientTransformation 实例。

另请参阅

功能优化器包装器 torchopt.FuncOptimizer


经典优化器

AdaDelta(params[, lr, rho, eps, weight_decay])

经典的AdaDelta优化器。

Adadelta

AdaDelta 的别名

AdaGrad(params[, lr, lr_decay, ...])

经典的AdaGrad优化器。

Adagrad

AdaGrad 的别名

Adam(params[, lr, betas, eps, weight_decay, ...])

经典的Adam优化器。

AdamW(params[, lr, betas, eps, ...])

经典的AdamW优化器。

AdaMax(params[, lr, betas, eps, weight_decay])

经典的AdaMax优化器。

Adamax

AdaMax 的别名

RAdam(params[, lr, betas, eps, weight_decay])

经典的RAdam优化器。

RMSProp(params[, lr, alpha, eps, ...])

经典的RMSProp优化器。

SGD(params, lr[, momentum, weight_decay, ...])

经典的SGD优化器。

经典AdaDelta优化器

class torchopt.AdaDelta(params, lr=1.0, rho=0.9, eps=1e-06, weight_decay=0.0)[source]

基础:Optimizer

经典的AdaDelta优化器。

另请参阅

  • 功能性的AdaDelta优化器:torchopt.adadelta()

  • 可微分的元AdaDelta优化器:torchopt.MetaAdaDetla.

初始化AdaDelta优化器。

Parameters:
  • params (可迭代的 Tensor) – 一个包含 torch.Tensor 的可迭代对象。指定应该优化的张量。

  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认: 1e-3)

  • rho (float, 可选) – 用于计算梯度及其平方的运行平均值的系数。 (默认值: 0.9)

  • eps (float, 可选) – 应用于平方根的一个小常数(如AdaDelta论文中所述),以避免在重新缩放时除以零。 (默认: 1e-6)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

经典AdaGrad优化器

class torchopt.AdaGrad(params, lr=0.01, lr_decay=0.0, weight_decay=0.0, initial_accumulator_value=0.0, eps=1e-10, *, maximize=False)[source]

基础:Optimizer

经典的AdaGrad优化器。

另请参阅

初始化AdaGrad优化器。

Parameters:
  • params (可迭代的 Tensor) – 一个包含 torch.Tensor 的可迭代对象。指定应该优化的张量。

  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认值: 1e-2)

  • lr_decay (float, 可选) – 学习率衰减。(默认: 0.0)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • initial_accumulator_value (float, 可选) – 累加器的初始值。 (默认: 0.0)

  • eps (float, optional) – 在平方根外应用于分母的小常数(如Adam论文中所述),以避免在重新缩放时除以零。 (default: 1e-10)

  • maximize (bool, optional) – 根据目标最大化参数,而不是最小化。(默认: False)

经典Adam优化器

class torchopt.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, *, eps_root=0.0, maximize=False, use_accelerated_op=False)[source]

基础:Optimizer

经典的Adam优化器。

另请参阅

初始化Adam优化器。

Parameters:
  • params (可迭代的 Tensor) – 一个包含 torch.Tensor 的可迭代对象。指定应该优化的张量。

  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认值: 1e-3)

  • betas (tuple of float, optional) – 用于计算梯度及其平方的运行平均值的系数。(默认值: (0.9, 0.999))

  • eps (float, optional) – 一个小的常数,应用于平方根外的分母(如Adam论文中所述),以避免在重新缩放时除以零。 (default: 1e-8)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • eps_root (float, optional) – 在平方根内的分母上应用的一个小常数(如RMSProp中所示),以避免在重新缩放时除以零。这在通过Adam计算(元)梯度时是必需的。(默认值:0.0

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建 带有标志 requires_grad=True,这个标志通常用于元学习 算法中。(默认值:False

  • maximize (bool, optional) – 根据目标最大化参数,而不是最小化。(默认: False)

  • use_accelerated_op (bool, 可选) – 如果 True 使用我们实现的融合操作符。 (默认: False)

经典AdamW优化器

class torchopt.AdamW(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.01, *, eps_root=0.0, mask=None, maximize=False, use_accelerated_op=False)[source]

基础:Optimizer

经典的AdamW优化器。

另请参阅

初始化AdamW优化器。

Parameters:
  • params (可迭代的 Tensor) – 一个包含 torch.Tensor 的可迭代对象。指定应该优化的张量。

  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认值: 1e-3)

  • betas (tuple of float, optional) – 用于计算梯度及其平方的运行平均值的系数。(默认值: (0.9, 0.999))

  • eps (float, 可选) – 一个小的常数,应用于平方根外的分母(如Adam论文中所述),以避免在重新缩放时除以零。 (默认: 1e-8)

  • weight_decay (float, 可选) – 权重衰减正则化的强度。请注意, 这个权重衰减会与学习率相乘。这与PyTorch等其他框架一致,但与(Loshchilov et al., 2019)不同, 在后者中,权重衰减仅与“计划乘数”相乘,而不与基础学习率相乘。(默认值:1e-2

  • eps_root (float, optional) – 在平方根内的分母上应用的一个小常数(如在RMSProp中),以避免在重新缩放时除以零。这在例如通过Adam计算(元)梯度时是必需的。(默认: 0.0)

  • mask (tree of Tensor, callable, or None, optional) – 一个与参数pytree结构相同(或为其前缀)的树,或者一个根据参数/更新返回此类pytree的函数。叶子节点应为布尔值, True 表示您希望应用权重衰减的叶子/子树,而 False 表示您希望跳过的叶子/子树。请注意,Adam梯度变换将应用于所有参数。(默认值:None

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建 带有标志 requires_grad=True,这个标志通常用于元学习 算法中。(默认:False

  • maximize (bool, optional) – 根据目标最大化参数,而不是最小化。(默认: False)

  • use_accelerated_op (bool, 可选) – 如果 True 使用我们实现的融合操作符。 (默认: False)

经典AdaMax优化器

class torchopt.AdaMax(params, lr=0.002, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0)[source]

基础:Optimizer

经典的AdaMax优化器。

另请参阅

初始化AdaMax优化器。

Parameters:
  • params (可迭代的 Tensor) – 一个包含 torch.Tensor 的可迭代对象。指定应该优化的张量。

  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认: 1e-3)

  • betas (tuple of float, optional) – 用于计算梯度及其平方的运行平均值的系数。(默认值: (0.9, 0.999))

  • eps (float, 可选) – 一个小的常数,应用于平方根(如AdaMax论文中所述),以避免在重新缩放时除以零。 (默认: 1e-6)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

经典RAdam优化器

class torchopt.RAdam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0)[source]

基础:Optimizer

经典的RAdam优化器。

另请参阅

初始化RAdam优化器。

Parameters:
  • params (可迭代的 Tensor) – 一个包含 torch.Tensor 的可迭代对象。指定应该优化的张量。

  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认: 1e-3)

  • betas (tuple of float, optional) – 用于计算梯度及其平方的运行平均值的系数。(默认值: (0.9, 0.999))

  • eps (float, 可选) – 一个应用于平方根的小常数(如RAdam论文中所述),以避免在重新缩放时除以零。 (默认: 1e-6)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

经典RMSProp优化器

class torchopt.RMSProp(params, lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0.0, momentum=0.0, centered=False, *, initial_scale=0.0, nesterov=False, maximize=False)[source]

基础:Optimizer

经典的RMSProp优化器。

另请参阅

初始化RMSProp优化器。

Parameters:
  • params (可迭代的 Tensor) – 一个包含 torch.Tensor 的可迭代对象。指定应该优化的张量。

  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认值: 1e-2)

  • alpha (float, 可选) – 平滑常数,用于跟踪先前梯度幅度的衰减。(默认: 0.99)

  • eps (float, 可选) – 一个小的数值常数,用于在重新缩放时避免除以零。(默认: 1e-8)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • 动量 (浮点数, 可选) – 动量项使用的衰减率。当设置为0.0时不使用动量。(默认值: 0.0)

  • 居中 (bool, 可选) – 如果 True,使用过去梯度的方差来重新缩放最新的梯度。(默认: False)

  • initial_scale (float, optional) – 用于初始化跟踪先前更新幅度的累加器。PyTorch 使用 0.0,TensorFlow 1.x 使用 1.0。在复现论文结果时,请验证作者使用的值。 (默认值: 0.0)

  • nesterov (bool, 可选) – 是否使用Nesterov动量。(默认: False)

  • maximize (bool, optional) – 根据目标最大化参数,而不是最小化。(默认: False)

经典SGD优化器

class torchopt.SGD(params, lr, momentum=0.0, weight_decay=0.0, dampening=0.0, nesterov=False, maximize=False)[source]

基础:Optimizer

经典的SGD优化器。

另请参阅

初始化SGD优化器。

Parameters:
  • params (可迭代的 Tensor) – 一个包含 torch.Tensor 的可迭代对象。指定应该优化的张量。

  • lr (floatcallable) – 这是一个固定的全局缩放因子或学习率调度器。

  • 动量 (float, 可选) – 动量项使用的衰减率。当设置为0.0时不使用动量。(默认值: 0.0)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • dampening (float, 可选) – 动量的阻尼系数。(默认: 0.0)

  • nesterov (bool, 可选) – 是否使用Nesterov动量。(默认: False)

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建 带有标志 requires_grad=True,这个标志通常用于元学习 算法中。(默认:False

  • maximize (bool, optional) – 根据目标最大化参数,而不是最小化。(默认: False)


可微分的元优化器

MetaAdaDelta(module[, lr, rho, eps, ...])

可微分的AdaDelta优化器。

MetaAdadelta

MetaAdaDelta 的别名

MetaAdaGrad(module[, lr, lr_decay, ...])

可微分的AdaGrad优化器。

MetaAdagrad

MetaAdaGrad 的别名

MetaAdam(module[, lr, betas, eps, ...])

可微分的Adam优化器。

MetaAdamW(module[, lr, betas, eps, ...])

可微分的AdamW优化器。

MetaAdaMax(module[, lr, betas, eps, ...])

可微分的AdaMax优化器。

MetaAdamax

MetaAdaMax 的别名

MetaRAdam(module[, lr, betas, eps, ...])

可微分的RAdam优化器。

MetaRMSProp(module[, lr, alpha, eps, ...])

可微分的RMSProp优化器。

MetaSGD(module, lr[, momentum, ...])

可微分的随机梯度下降优化器。

可微分的元AdaDelta优化器

class torchopt.MetaAdaDelta(module, lr=1.0, rho=0.9, eps=1e-06, weight_decay=0.0, *, moment_requires_grad=True)[source]

基础类: MetaOptimizer

可微分的AdaDelta优化器。

另请参阅

  • 功能性的AdaDelta优化器:torchopt.adadetla()

  • 经典的AdaDelta优化器:torchopt.Adadelta

初始化元AdaDelta优化器。

Parameters:
  • module (nn.Module) – 一个需要优化参数的网络。

  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认: 1e-3)

  • rho (float, 可选) – 用于计算梯度及其平方的运行平均值的系数。 (默认值: 0.9)

  • eps (float, 可选) – 应用于平方根的一个小常数(如AdaDelta论文中所述),以避免在重新缩放时除以零。 (默认: 1e-6)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建 带有标志 requires_grad=True,这个标志通常用于元学习 算法中。(默认:False

可微分的Meta-AdaGrad优化器

class torchopt.MetaAdaGrad(module, lr=0.01, lr_decay=0.0, weight_decay=0.0, initial_accumulator_value=0.0, eps=1e-10, *, maximize=False)[source]

基础类: MetaOptimizer

可微分的AdaGrad优化器。

另请参阅

  • 功能性的AdaGrad优化器:torchopt.adagrad()

  • 经典的AdaGrad优化器:torchopt.Adagrad

初始化元AdaGrad优化器。

Parameters:
  • module (nn.Module) – 一个需要优化参数的网络。

  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认值: 1e-2)

  • lr_decay (float, 可选) – 学习率衰减。(默认: 0.0)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • initial_accumulator_value (float, 可选) – 累加器的初始值。 (默认: 0.0)

  • eps (float, optional) – 在平方根外应用于分母的小常数(如Adam论文中所述),以避免在重新缩放时除以零。 (default: 1e-10)

  • maximize (bool, optional) – 根据目标最大化参数,而不是最小化。(默认: False)

可微分的Meta-Adam优化器

class torchopt.MetaAdam(module, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, *, eps_root=0.0, moment_requires_grad=True, maximize=False, use_accelerated_op=False)[source]

基础类: MetaOptimizer

可微分的Adam优化器。

另请参阅

初始化元-Adam优化器。

Parameters:
  • module (nn.Module) – 一个需要优化参数的网络。

  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认值: 1e-3)

  • betas (tuple of float, optional) – 用于计算梯度及其平方的运行平均值的系数。(默认值: (0.9, 0.999))

  • eps (float, 可选) – 一个小的常数,应用于平方根外的分母(如Adam论文中所述),以避免在重新缩放时除以零。 (默认: 1e-8)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • eps_root (float, optional) – 在平方根内的分母上应用的一个小常数(如在RMSProp中),以避免在重新缩放时除以零。这在例如通过Adam计算(元)梯度时是必需的。(默认: 0.0)

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建 带有标志 requires_grad=True,这个标志通常用于元学习 算法中。(默认:False

  • maximize (bool, optional) – 根据目标最大化参数,而不是最小化。(默认: False)

  • use_accelerated_op (bool, 可选) – 如果 True 使用我们实现的融合操作符。 (默认: False)

可微分的Meta-AdamW优化器

class torchopt.MetaAdamW(module, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.01, *, eps_root=0.0, mask=None, moment_requires_grad=False, maximize=False, use_accelerated_op=False)[source]

基础类: MetaOptimizer

可微分的AdamW优化器。

另请参阅

初始化元-AdamW优化器。

Parameters:
  • module (nn.Module) – 一个需要优化参数的网络。

  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认值: 1e-3)

  • betas (tuple of float, optional) – 用于计算梯度及其平方的运行平均值的系数。(默认值: (0.9, 0.999))

  • eps (float, 可选) – 一个小的常数,应用于平方根外的分母(如Adam论文中所述),以避免在重新缩放时除以零。 (默认: 1e-8)

  • weight_decay (float, 可选) – 权重衰减正则化的强度。请注意, 这个权重衰减会与学习率相乘。这与PyTorch等其他框架一致,但与(Loshchilov et al., 2019)不同, 在后者中,权重衰减仅与“计划乘数”相乘,而不与基础学习率相乘。(默认值:1e-2

  • eps_root (float, optional) – 在平方根内的分母上应用的一个小常数(如在RMSProp中),以避免在重新缩放时除以零。这在例如通过Adam计算(元)梯度时是必需的。(默认: 0.0)

  • mask (tree of Tensor, callable, or None, optional) – 一个与参数pytree结构相同(或为其前缀)的树,或者一个根据参数/更新返回此类pytree的函数。叶子节点应为布尔值, True 表示您希望应用权重衰减的叶子/子树,而 False 表示您希望跳过的叶子/子树。请注意,Adam梯度变换将应用于所有参数。(默认值:None

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建 带有标志 requires_grad=True,这个标志通常用于元学习 算法中。(默认:False

  • maximize (bool, optional) – 根据目标最大化参数,而不是最小化。(默认: False)

  • use_accelerated_op (bool, 可选) – 如果 True 使用我们实现的融合操作符。 (默认: False)

可微分的Meta-AdaMax优化器

class torchopt.MetaAdaMax(module, lr=0.002, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, *, moment_requires_grad=True)[source]

基础类: MetaOptimizer

可微分的AdaMax优化器。

另请参阅

  • 功能性的AdaMax优化器:torchopt.adamax()

  • 经典的AdaMax优化器:torchopt.Adamax

初始化元AdaMax优化器。

Parameters:
  • module (nn.Module) – 一个需要优化参数的网络。

  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认: 1e-3)

  • betas (tuple of float, optional) – 用于计算梯度及其平方的运行平均值的系数。(默认值: (0.9, 0.999))

  • eps (float, 可选) – 一个小的常数,应用于平方根(如AdaMax论文中所述),以避免在重新缩放时除以零。 (默认: 1e-6)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建 带有标志 requires_grad=True,这个标志通常用于元学习 算法中。(默认:False

可微分的Meta-RAdam优化器

class torchopt.MetaRAdam(module, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, *, moment_requires_grad=True)[source]

基础类: MetaOptimizer

可微分的RAdam优化器。

另请参阅

  • 功能性的RAdam优化器:torchopt.radan()

  • 经典的RAdam优化器:torchopt.RAdam

初始化元RAdam优化器。

Parameters:
  • module (nn.Module) – 一个需要优化参数的网络。

  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认: 1e-3)

  • betas (tuple of float, optional) – 用于计算梯度及其平方的运行平均值的系数。(默认值: (0.9, 0.999))

  • eps (float, 可选) – 一个应用于平方根的小常数(如RAdam论文中所述),以避免在重新缩放时除以零。 (默认: 1e-6)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建 带有标志 requires_grad=True,这个标志通常用于元学习 算法中。(默认:False

可微分的Meta-RMSProp优化器

class torchopt.MetaRMSProp(module, lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0.0, momentum=0.0, centered=False, *, initial_scale=0.0, nesterov=False, maximize=False)[source]

基础类: MetaOptimizer

可微分的RMSProp优化器。

另请参阅

初始化元-RMSProp优化器。

Parameters:
  • module (nn.Module) – 一个需要优化参数的网络。

  • lr (floatcallable, 可选) – 这是一个固定的全局缩放因子或学习率调度器。(默认值: 1e-2)

  • alpha (float, 可选) – 平滑常数,用于跟踪先前梯度幅度的衰减。(默认: 0.99)

  • eps (float, 可选) – 一个小的数值常数,用于在重新缩放时避免除以零。(默认: 1e-8)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • 动量 (float, 可选) – 动量项使用的衰减率。当设置为0.0时不使用动量。(默认值: 0.0)

  • 居中 (bool, 可选) – 如果 True,使用过去梯度的方差来重新缩放最新的梯度。(默认: False)

  • initial_scale (float, optional) – 用于初始化跟踪先前更新幅度的累加器。PyTorch 使用 0.0,TensorFlow 1.x 使用 1.0。在复现论文结果时,请验证作者使用的值。 (默认值: 0.0)

  • nesterov (bool, 可选) – 是否使用Nesterov动量。(默认: False)

  • maximize (bool, optional) – 根据目标最大化参数,而不是最小化。(默认: False)

可微分的Meta-SGD优化器

class torchopt.MetaSGD(module, lr, momentum=0.0, weight_decay=0.0, dampening=0.0, nesterov=False, moment_requires_grad=True, maximize=False)[source]

基础类: MetaOptimizer

可微分的随机梯度下降优化器。

另请参阅

初始化元SGD优化器。

Parameters:
  • module (nn.Module) – 一个需要优化参数的网络。

  • lr (floatcallable) – 这是一个固定的全局缩放因子或学习率调度器。

  • 动量 (float, 可选) – 动量项使用的衰减率。当设置为0.0时不使用动量。(默认值: 0.0)

  • weight_decay (float, 可选) – 权重衰减,向参数添加L2惩罚。 (默认: 0.0)

  • dampening (float, 可选) – 动量的阻尼系数。(默认: 0.0)

  • nesterov (bool, 可选) – 是否使用Nesterov动量。(默认: False)

  • moment_requires_grad (bool, 可选) – 如果 True,动量将被创建 带有标志 requires_grad=True,这个标志通常用于元学习 算法中。(默认:False

  • maximize (bool, optional) – 根据目标最大化参数,而不是最小化。(默认: False)


隐函数微分

custom_root(optimality_fn, argnums[, ...])

返回一个装饰器,用于向根求解器添加隐式微分。

nn.ImplicitMetaGradientModule(*args, **kwargs)

可微分隐式元梯度模型的基类。

自定义求解器

torchopt.diff.implicit.custom_root(optimality_fn, argnums, has_aux=False, solve=None)[source]

返回一个装饰器,用于向根求解器添加隐式微分。

这个包装器应该用作装饰器:

def optimality_fn(optimal_params, ...):
    ...
    return residual

@custom_root(optimality_fn, argnums=argnums)
def solver_fn(params, arg1, arg2, ...):
    ...
    return optimal_params

optimal_params = solver_fn(init_params, ...)

optimality_fnsolver_fn 的第一个参数被保留为参数输入。 argnums 参数指的是 solver_fn 签名中变量的索引。 例如,设置 argnums=(1, 2) 将计算 optimal_params 相对于上述代码片段中的 arg1arg2 的梯度。请注意,除了第一个参数外,optimality_fn 的关键字参数应该是 solver_fn 的关键字参数的子集。 最佳实践是,``optimality_fn`` 应该具有与 ``solver_fn`` 相同的签名。

Parameters:
  • optimality_fn (callable) – 一个方程函数,optimality_fn(params, *args)。不变性是在解/根处optimality_fn(solution, *args) == 0

  • argnums (inttupleint) – 指定要计算梯度的参数。argnums 可以是一个整数或一个整数元组,对应于 solver_fn(params, *args) 函数的参数的从零开始的索引。参数 params 被包括在计数中,同时它被索引为 argnums=0

  • has_aux (bool, 可选) – 装饰的求解器函数是否返回辅助数据。 (默认: False)

  • solve (callable, optional) – 一个形式为 solve(matvec, b) 的线性求解器。 (默认: linear_solve.solve_normal_cg())

Return type:

Callable[[Callable[..., Union[Tensor, Sequence[Tensor], tuple[Union[Tensor, Sequence[Tensor]], Any]]]], Callable[..., Union[Tensor, Sequence[Tensor], tuple[Union[Tensor, Sequence[Tensor]], Any]]]]

Returns:

一个求解器函数装饰器,即 custom_root(optimality_fn)(solver_fn)

隐式元梯度模块

class torchopt.diff.implicit.nn.ImplicitMetaGradientModule(*args, **kwargs)[source]

基础类:MetaGradientModule

可微分隐式元梯度模型的基类。

初始化一个新的模块实例。

linear_solve: LinearSolver | None
classmethod __init_subclass__(linear_solve=None)[source]

验证并初始化子类。

Return type:

abstract solve(*input, **kwargs)[source]

解决内部优化问题。 :rtype: Any

警告

对于基于梯度的优化方法,参数输入应在torch.autograd.backward()函数中明确指定为参数inputs。否则,如果未提供,梯度将累积到用于计算目标输出的所有叶张量(包括元参数)中。或者,请使用torch.autograd.grad()代替。

示例

def solve(self, batch, labels):
    parameters = tuple(self.parameters())
    optimizer = torch.optim.Adam(parameters, lr=1e-3)
    with torch.enable_grad():
        for _ in range(100):
            loss = self.objective(batch, labels)
            optimizer.zero_grad()
            # Only update the `.grad` attribute for parameters
            # and leave the meta-parameters unchanged
            loss.backward(inputs=parameters)
            optimizer.step()
    return self
optimality(*input, **kwargs)[source]

计算最优性残差。

该方法代表在解决内部优化问题(solve())后,对最优参数的最优性残差,即:

module.solve(*input, **kwargs)
module.optimality(*input, **kwargs)  # -> 0

1. 对于基于梯度的优化,optimality() 函数是KKT条件, 通常是objective() 函数相对于模块参数的梯度 (而不是元参数)。如果未实现此方法,它将自动从objective() 函数的梯度中推导出来。

\[\text{optimality residual} = \nabla_{\boldsymbol{x}} f (\boldsymbol{x}, \boldsymbol{\theta}) \to \boldsymbol{0}\]

其中 \(\boldsymbol{x}\) 是模块参数的联合向量, \(\boldsymbol{\theta}\) 是元参数的联合向量。

参考文献

2. 对于定点迭代,optimality() 函数可以是参数在迭代之间的残差,即:

\[\text{optimality residual} = f (\boldsymbol{x}, \boldsymbol{\theta}) - \boldsymbol{x} \to \boldsymbol{0}\]

其中 \(\boldsymbol{x}\) 是模块参数的联合向量, \(\boldsymbol{\theta}\) 是元参数的联合向量。

Return type:

元组[张量, ...]

Returns:

一个张量元组,表示在解决内部优化问题后,最优参数的最优性残差。返回的张量应与tuple(self.parameters())的输出相对应。

objective(*input, **kwargs)[source]

计算目标函数值。

此方法用于计算optimality()(如果未实现)。 否则,此方法是可选的。

Return type:

Tensor

Returns:

一个标量张量 (dim=0),目标函数值。

__abstractmethods__ = frozenset({'solve'})
__annotations__ = {'__call__': 'Callable[..., Any]', '_backward_hooks': 'Dict[int, Callable]', '_backward_pre_hooks': 'Dict[int, Callable]', '_buffers': 'Dict[str, Optional[Tensor]]', '_compiled_call_impl': 'Optional[Callable]', '_custom_objective': 'bool', '_custom_optimality': 'bool', '_forward_hooks': 'Dict[int, Callable]', '_forward_hooks_always_called': 'Dict[int, bool]', '_forward_hooks_with_kwargs': 'Dict[int, bool]', '_forward_pre_hooks': 'Dict[int, Callable]', '_forward_pre_hooks_with_kwargs': 'Dict[int, bool]', '_is_full_backward_hook': 'Optional[bool]', '_load_state_dict_post_hooks': 'Dict[int, Callable]', '_load_state_dict_pre_hooks': 'Dict[int, Callable]', '_meta_inputs': 'MetaInputsContainer', '_meta_modules': 'dict[str, nn.Module | None]', '_meta_parameters': 'TensorContainer', '_modules': "Dict[str, Optional['Module']]", '_non_persistent_buffers_set': 'Set[str]', '_parameters': 'Dict[str, Optional[Parameter]]', '_state_dict_hooks': 'Dict[int, Callable]', '_state_dict_pre_hooks': 'Dict[int, Callable]', '_version': 'int', 'call_super_init': 'bool', 'dump_patches': 'bool', 'forward': 'Callable[..., Any]', 'linear_solve': 'LinearSolver | None', 'training': 'bool'}

线性系统求解器

solve_cg(**kwargs)

返回一个求解器函数,用于使用共轭梯度法求解 A x = b

solve_normal_cg(**kwargs)

返回一个求解器函数,用于使用共轭梯度法求解 A^T A x = A^T b

solve_inv(**kwargs)

返回一个求解器函数,用于使用矩阵逆来求解 A x = b

间接求解器

torchopt.linear_solve.solve_cg(**kwargs)[source]

返回一个求解器函数,用于使用共轭梯度法求解 A x = b

这假设A是一个厄米特正定矩阵。

Parameters:
  • ridge (floatNone, 可选) – 可选的岭正则化。如果提供,则求解方程 A x + ridge x = b。 (默认: None)

  • init (Tensor, tree of Tensor, or None, optional) – 可选的初始化方法,用于共轭梯度。如果 None,则使用零初始化。(默认值:None

  • **kwargs (Any) – 共轭梯度求解器的附加关键字参数 torchopt.linalg.cg().

Return type:

Callable[[Callable[[Union[Tensor, Tuple[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]], ...], List[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Dict[Any, Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Deque[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], CustomTreeNode[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]]]], Union[Tensor, Tuple[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]], ...], List[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Dict[Any, Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Deque[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], CustomTreeNode[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]]]], Union[Tensor, Tuple[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]], ...], List[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Dict[Any, Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Deque[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], CustomTreeNode[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]]]], Union[Tensor, Tuple[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]], ...], List[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Dict[Any, Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Deque[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], CustomTreeNode[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]]]]

Returns:

一个具有签名 (matvec, b) -> x 的求解函数,用于使用共轭梯度法求解 A x = b,其中 matvec(v) = A v

另请参阅

共轭梯度迭代 torchopt.linalg.cg().

示例

>>> A = {'a': torch.eye(5, 5), 'b': torch.eye(3, 3)}
>>> x = {'a': torch.randn(5), 'b': torch.randn(3)}
>>> def matvec(x: TensorTree) -> TensorTree:
...     return {'a': A['a'] @ x['a'], 'b': A['b'] @ x['b']}
>>> b = matvec(x)
>>> solver = solve_cg(init={'a': torch.zeros(5), 'b': torch.zeros(3)})
>>> x_hat = solver(matvec, b)
>>> assert torch.allclose(x_hat['a'], x['a']) and torch.allclose(x_hat['b'], x['b'])
torchopt.linear_solve.solve_normal_cg(**kwargs)[source]

返回一个求解器函数,用于使用共轭梯度法求解 A^T A x = A^T b

这可以用于在A不是厄米特正定时,使用共轭梯度法求解A x = b

Parameters:
  • ridge (floatNone, 可选) – 可选的岭正则化。如果提供,则求解方程 A^T A x + ridge x = A^T b。 (默认: None)

  • init (Tensor, tree of Tensor, or None, optional) – 可选的初始化方法,用于共轭梯度。如果 None,则使用零初始化。(默认值:None

  • **kwargs (Any) – 共轭梯度求解器的附加关键字参数 torchopt.linalg.cg().

Return type:

Callable[[Callable[[Union[Tensor, Tuple[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]], ...], List[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Dict[Any, Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Deque[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], CustomTreeNode[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]]]], Union[Tensor, Tuple[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]], ...], List[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Dict[Any, Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Deque[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], CustomTreeNode[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]]]], Union[Tensor, Tuple[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]], ...], List[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Dict[Any, Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Deque[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], CustomTreeNode[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]]]], Union[Tensor, Tuple[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]], ...], List[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Dict[Any, Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Deque[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], CustomTreeNode[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]]]]

Returns:

一个具有签名(matvec, b) -> x的求解函数,用于使用共轭梯度法求解A^T A x = A^T b,其中matvec(v) = A v

另请参阅

共轭梯度迭代 torchopt.linalg.cg().

示例

>>> A = {'a': torch.randn(5, 5), 'b': torch.randn(3, 3)}
>>> x = {'a': torch.randn(5), 'b': torch.randn(3)}
>>> def matvec(x: TensorTree) -> TensorTree:
...     return {'a': A['a'] @ x['a'], 'b': A['b'] @ x['b']}
>>> b = matvec(x)
>>> solver = solve_normal_cg(init={'a': torch.zeros(5), 'b': torch.zeros(3)})
>>> x_hat = solver(matvec, b)
>>> assert torch.allclose(x_hat['a'], x['a']) and torch.allclose(x_hat['b'], x['b'])
torchopt.linear_solve.solve_inv(**kwargs)[source]

返回一个求解器函数,用于使用矩阵逆解A x = b

如果 ns = False,这假设矩阵 A 是一个常数矩阵,并将其存储在内存中。

Parameters:
  • ridge (floatNone, 可选) – 可选的岭正则化。如果提供,则求解方程 A x + ridge x = b。 (默认: None)

  • ns (bool, 可选) – 是否使用Neumann Series矩阵逆近似。 如果 False,则将矩阵 A 具体化到内存中并使用 torch.linalg.solve() 代替。(默认值:False

  • **kwargs (Any) – 用于Neumann级数矩阵逆近似求解器 torchopt.linalg.ns() 的额外关键字参数。

Return type:

Callable[[Callable[[Union[Tensor, Tuple[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]], ...], List[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Dict[Any, Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Deque[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], CustomTreeNode[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]]]], Union[Tensor, Tuple[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]], ...], List[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Dict[Any, Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Deque[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], CustomTreeNode[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]]]], Union[Tensor, Tuple[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]], ...], List[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Dict[Any, Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Deque[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], CustomTreeNode[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]]]], Union[Tensor, Tuple[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]], ...], List[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Dict[Any, Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Deque[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], CustomTreeNode[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]]]]

Returns:

一个具有签名(matvec, b) -> x的求解函数,使用矩阵逆来求解A x = b,其中matvec(v) = A v

另请参阅

诺伊曼级数矩阵逆近似 torchopt.linalg.ns().

示例

>>> A = {'a': torch.eye(5, 5), 'b': torch.eye(3, 3)}
>>> x = {'a': torch.randn(5), 'b': torch.randn(3)}
>>> def matvec(x: TensorTree) -> TensorTree:
...     return {'a': A['a'] @ x['a'], 'b': A['b'] @ x['b']}
>>> b = matvec(x)
>>> solver = solve_inv(ns=True, maxiter=10)
>>> x_hat = solver(matvec, b)
>>> assert torch.allclose(x_hat['a'], x['a']) and torch.allclose(x_hat['b'], x['b'])

零阶微分

zero_order(distribution[, method, argnums, ...])

返回一个用于应用零阶微分的装饰器。

nn.ZeroOrderGradientModule(*args, **kwargs)

零阶梯度模型的基类。

装饰器

torchopt.diff.zero_order.zero_order(distribution, method='naive', argnums=(0,), num_samples=1, sigma=1.0)[source]

返回一个用于应用零阶微分的装饰器。

Parameters:
  • 分布 (可调用的可采样的) – 一个可采样的对象,具有方法 samplable.sample(sample_shape) 或一个函数,该函数以形状为输入并 返回一个形状化的样本批次。这用于从给定的分布中采样扰动。分布应该是球对称的。

  • 方法 (str, 可选) – 使用的算法。当前支持的算法是 'naive', 'forward', 和 'antithetic'. (默认: 'naive')

  • argnums (inttupleint, 可选) – 指定要计算梯度的参数。(默认: 0)

  • num_samples (int, optional) – 获取平均估计梯度所需的样本数量。 (默认: 1)

  • sigma (float, 可选) – 扰动的标准差。 (默认: 1.0)

Return type:

Callable[[Callable[..., Tensor]], Callable[..., Tensor]]

Returns:

一个函数装饰器,用于启用零阶梯度估计。

零阶梯度模块

class torchopt.diff.zero_order.nn.ZeroOrderGradientModule(*args, **kwargs)[source]

基础类:Module, Samplable

零阶梯度模型的基础类。

初始化内部模块状态,由nn.Module和ScriptModule共享。

classmethod __init_subclass__(method='naive', num_samples=1, sigma=1.0)[source]

验证并初始化子类。

Return type:

abstract forward(*args, **kwargs)[source]

执行模型的前向传递。

Return type:

Tensor

abstract sample(sample_shape=())[source]

如果分布参数是批处理的,则生成一个sample_shape形状的样本或sample_shape形状的样本批次。

Return type:

Union[Tensor, Sequence[Union[Tensor, float, int, bool]]]

__abstractmethods__ = frozenset({'forward', 'sample'})
__annotations__ = {'__call__': 'Callable[..., Any]', '_backward_hooks': 'Dict[int, Callable]', '_backward_pre_hooks': 'Dict[int, Callable]', '_buffers': 'Dict[str, Optional[Tensor]]', '_compiled_call_impl': 'Optional[Callable]', '_forward_hooks': 'Dict[int, Callable]', '_forward_hooks_always_called': 'Dict[int, bool]', '_forward_hooks_with_kwargs': 'Dict[int, bool]', '_forward_pre_hooks': 'Dict[int, Callable]', '_forward_pre_hooks_with_kwargs': 'Dict[int, bool]', '_is_full_backward_hook': 'Optional[bool]', '_load_state_dict_post_hooks': 'Dict[int, Callable]', '_load_state_dict_pre_hooks': 'Dict[int, Callable]', '_modules': "Dict[str, Optional['Module']]", '_non_persistent_buffers_set': 'Set[str]', '_parameters': 'Dict[str, Optional[Parameter]]', '_state_dict_hooks': 'Dict[int, Callable]', '_state_dict_pre_hooks': 'Dict[int, Callable]', '_version': 'int', 'call_super_init': 'bool', 'dump_patches': 'bool', 'forward': 'Callable[..., Any]', 'training': 'bool'}
__parameters__ = ()
__subclasshook__()

抽象类可以重写此方法以自定义 issubclass()。

这是由abc.ABCMeta.__subclasscheck__()早期调用的。 它应该返回True、False或NotImplemented。如果返回 NotImplemented,则使用正常算法。否则,它 会覆盖正常算法(并且结果会被缓存)。


优化器钩子

register_hook(hook)

无状态的身份转换,保持输入梯度不变。

zero_nan_hook(g)

nan 替换为零。

nan_to_num_hook([nan, posinf, neginf])

返回一个nan到num的钩子,用于将nan / +inf / -inf替换为给定的数字。

钩子

torchopt.hook.register_hook(hook)[source]

无状态的身份转换,保持输入梯度不变。

此函数通过梯度更新保持不变。

Return type:

GradientTransformation

Returns:

一个 (init_fn, update_fn) 元组。

torchopt.hook.zero_nan_hook(g)[source]

nan替换为零。

Return type:

Tensor

torchopt.hook.nan_to_num_hook(nan=0.0, posinf=None, neginf=None)[source]

返回一个nan到num的钩子,用给定的数字替换nan / +inf / -inf

Return type:

Callable[[Tensor], Tensor]


梯度变换

clip_grad_norm(max_norm[, norm_type, ...])

裁剪参数迭代器的梯度范数。

nan_to_num([nan, posinf, neginf])

将更新替换为给定的数字 nan / +inf / -inf

转换

torchopt.clip_grad_norm(max_norm, norm_type=2.0, error_if_nonfinite=False)[source]

裁剪参数可迭代对象的梯度范数。

Parameters:
  • max_norm (float) – 更新中每个元素的最大绝对值。

  • norm_type (float, 可选) – 使用的p-范数类型。可以是'inf'表示无穷范数。 (默认: 2.0)

  • error_if_nonfinite (bool, 可选) – 如果 True,当来自 updates 的梯度的总范数为 naninf-inf 时,将抛出错误。 (默认: False)

Return type:

GradientTransformation

Returns:

一个 (init_fn, update_fn) 元组。

torchopt.nan_to_num(nan=0.0, posinf=None, neginf=None)[source]

将更新替换为给定数字的nan / +inf / -inf值。

Return type:

GradientTransformation

Returns:

一个 (init_fn, update_fn) 元组。

优化器调度

linear_schedule(init_value, end_value, ...)

为了方便起见,将多项式调度别名为线性调度。

polynomial_schedule(init_value, end_value, ...)

构建一个从初始值到结束值具有多项式过渡的调度。

日程安排

torchopt.schedule.linear_schedule(init_value, end_value, transition_steps, transition_begin=0)[source]

为了方便起见,将多项式调度别名为线性调度。

Return type:

Callable[[Union[Tensor, float, int, bool]], Union[Tensor, float, int, bool]]

torchopt.schedule.polynomial_schedule(init_value, end_value, power, transition_steps, transition_begin=0)[source]

构建一个从初始值到结束值具有多项式过渡的调度。

Parameters:
  • init_value (floatTensor) – 要退火的标量的初始值。

  • end_value (floatTensor) – 要退火的标量的结束值。

  • power (floatTensor) – 用于从 init 过渡到 end 的多项式的幂。

  • transition_steps (int) – 退火过程发生的步数,标量从 transition_begin 步开始变化,并在 transition_begin + transition_steps 步完成过渡。如果 transition_steps <= 0,则整个退火过程被禁用,并且值固定在 init_value

  • transition_begin (int, 可选) – 必须为正数。在多少步之后开始退火(在此之前,标量值固定在init_value)。 (默认值: 0)

Returns:

一个将步数映射到值的函数。

Return type:

日程安排

应用参数更新

apply_updates(params, updates, *[, inplace])

对相应的参数应用更新。

应用更新

torchopt.apply_updates(params, updates, *, inplace=True)[source]

对相应的参数应用更新。

这是一个实用函数,用于对一组参数应用更新,然后将更新后的参数返回给调用者。例如,更新可能是由一系列GradientTransformations转换的梯度。为了方便起见,这个函数被公开,但它只是添加更新和参数;你也可以手动将更新应用到参数上,使用tree_map()(例如,如果你想在应用更新之前以自定义方式操作更新)。

Parameters:
  • params (tree of Tensor) – 参数的树结构。

  • 更新 (张量) – 更新的树,树结构和叶子节点的形状必须与params匹配。

  • inplace (bool, 可选) – 如果为True,将以原地方式更新参数。 (默认: True)

Return type:

参数

Returns:

更新后的参数,与params具有相同的结构、形状和类型。

组合优化器

chain(*transformations)

应用一系列可链式更新的转换。

torchopt.combine.chain(*transformations)[source]

应用一系列可链式更新的转换。

给定一系列可链接的转换,chain() 返回一个 init_fn(),它通过连接各个转换的状态来构建一个 state,并返回一个 update_fn(),该函数将更新转换链接起来,将适当的状态传递给每个转换。

Parameters:

*transformations (iterable of GradientTransformation) – 一系列可链式调用的 (init_fn, update_fn) 元组。

Return type:

GradientTransformation

Returns:

一个单一的 (init_fn, update_fn) 元组。

分布式工具

初始化和同步

auto_init_rpc([worker_init_fn, ...])

返回一个装饰器,用于自动在被装饰的函数上初始化RPC。

barrier([worker_names])

同步本地和远程的RPC进程。

torchopt.distributed.auto_init_rpc(worker_init_fn=None, worker_name_format=<function default_worker_name_format>, *, backend=None, rpc_backend_options=None)[source]

返回一个装饰器,用于自动在被装饰的函数上初始化RPC。

Return type:

Callable[[TypeVar(F, bound= Callable[..., Any])], TypeVar(F, bound= Callable[..., Any])]

torchopt.distributed.barrier(worker_names=None)[source]

同步本地和远程RPC进程。

这将阻塞,直到在worker_names下指定的所有本地和远程RPC进程到达此方法,以等待所有未完成的工作完成。

Parameters:

worker_names (可迭代的 strNone, 可选) – 要同步的工作者集合。 如果 None,则同步所有工作者。(默认值:None

Return type:

处理组信息

get_world_info()

获取世界信息。

get_world_rank()

获取当前工作者的全球排名。

get_rank()

获取当前工作者的全球排名。

get_world_size()

获取世界大小。

get_local_rank()

获取当前节点上当前工作进程的本地排名。

get_local_world_size()

获取当前节点上的本地世界大小。

get_worker_id([id])

从给定的id中获取工作id。

torchopt.distributed.get_world_info()[source]

获取世界信息。

Return type:

WorldInfo

torchopt.distributed.get_world_rank()[source]

获取当前工作者的全球排名。

Return type:

int

torchopt.distributed.get_rank()

获取当前工作者的全球排名。

Return type:

int

torchopt.distributed.get_world_size()[source]

获取世界大小。

Return type:

int

torchopt.distributed.get_local_rank()[source]

获取当前节点上当前工作者的本地排名。

Return type:

int

torchopt.distributed.get_local_world_size()[source]

获取当前节点上的本地世界大小。

Return type:

int

torchopt.distributed.get_worker_id(id=None)[source]

从给定的id中获取工作者的id。

Return type:

int

工人选择

on_rank(*ranks)

返回一个装饰器,用于标记一个函数仅在给定的rank上执行。

not_on_rank(*ranks)

返回一个装饰器,用于标记一个函数仅在非给定等级上执行。

rank_zero_only(func)

返回一个装饰器,用于标记一个函数仅在rank zero上执行。

rank_non_zero_only(func)

返回一个装饰器,用于标记一个函数仅在非零排名上执行。

torchopt.distributed.on_rank(*ranks)[source]

返回一个装饰器,用于标记仅在给定等级上执行的函数。

Return type:

Callable[[TypeVar(F, bound= Callable[..., Any])], TypeVar(F, bound= Callable[..., Any])]

torchopt.distributed.not_on_rank(*ranks)[source]

返回一个装饰器,用于标记一个函数仅在未给定的等级上执行。

Return type:

Callable[[TypeVar(F, bound= Callable[..., Any])], TypeVar(F, bound= Callable[..., Any])]

torchopt.distributed.rank_zero_only(func)[source]

返回一个装饰器,用于标记一个函数仅在排名为零时执行。

Return type:

TypeVar(F, bound= Callable[..., Any])

torchopt.distributed.rank_non_zero_only(func)[source]

返回一个装饰器,用于标记一个函数仅在非零等级上执行。

Return type:

TypeVar(F, bound= Callable[..., Any])

远程过程调用 (RPC)

remote_async_call(func, *[, args, kwargs, ...])

在远程工作者上异步执行RPC并返回当前工作者上的一个torch.Future实例。

remote_sync_call(func, *[, args, kwargs, ...])

在远程工作者上同步执行RPC并将结果返回给当前工作者。

torchopt.distributed.remote_async_call(func, *, args=None, kwargs=None, partitioner=None, reducer=None, timeout=-1.0)[source]

在远程工作者上异步执行RPC并返回当前工作者处的torch.Future实例。

Parameters:
  • func (callable) – 要调用的函数。

  • args (tuple of object or None, optional) – 传递给函数的参数。 (默认值: None)

  • kwargs (dict[str, object] 或 None, 可选) – 传递给函数的关键字参数。 (默认: None)

  • partitioner (int, str, 或 callable, 可选) – 一个将参数分区到多个工作节点的分区器。(默认值:batch_partitioner()

  • reducer (callableNone, 可选) – 一个用于减少多个工作者结果的reducer。如果 None,则不减少结果。(默认值:None

  • timeout (float, 可选) – RPC调用的超时时间。 (默认: rpc.api.UNSET_RPC_TIMEOUT)

Return type:

Union[Future[list[TypeVar(T)]], Future[TypeVar(U)]]

Returns:

一个torch.Future实例用于结果。结果在当前工作节点。

torchopt.distributed.remote_sync_call(func, *, args=None, kwargs=None, partitioner=None, reducer=None, timeout=-1.0)[source]

在远程工作者上同步执行RPC并将结果返回给当前工作者。

Parameters:
  • func (callable) – 要调用的函数。

  • args (tuple of object or None, optional) – 传递给函数的参数。 (默认值: None)

  • kwargs (dict[str, object] 或 None, 可选) – 传递给函数的关键字参数。 (默认值: None)

  • partitioner (int, str, 或 callable, 可选) – 一个将参数分区到多个工作节点的分区器。(默认: batch_partitioner())

  • reducer (callableNone, 可选) – 一个用于减少多个工作者结果的reducer。如果 None,则不减少结果。(默认值:None

  • timeout (float, 可选) – RPC调用的超时时间。 (默认: rpc.api.UNSET_RPC_TIMEOUT)

Return type:

Union[list[TypeVar(T)], TypeVar(U)]

Returns:

RPC调用的结果。结果在当前工作节点。

预定义的分区和归约器

dim_partitioner([dim, exclusive, keepdim, ...])

沿给定维度对一批输入进行分区。

batch_partitioner

分区器类,用于沿给定维度对一批输入进行分区。

mean_reducer(results)

通过平均结果来减少结果。

sum_reducer(results)

通过求和来减少结果。

torchopt.distributed.dim_partitioner(dim=0, *, exclusive=False, keepdim=True, workers=None)[source]

沿给定维度对一批输入进行分区。

argskwargs中的所有张量将沿着维度dim进行分区,而非张量值将被广播到各个分区。

Parameters:
  • dim (int, 可选) – 要分区的维度。(默认: 0)

  • exclusive (bool, 可选) – 是否独占分区批次。(默认: False) 如果 True,批次将被分区为 batch_size 个分区,其中 batch_size 是批次在给定维度上的大小。每个批次样本 将被分配给一个单独的 RPC 调用。 如果 False,批次将被分区为 min(batch_size, num_workers) 个分区,其中 num_workers 是世界中的工作进程数。当 batch_size > num_workers 时,单个 RPC 调用中可以有多个批次样本。

  • keepdim (bool, 可选) – 是否保留分区维度。(默认: False) 如果 True,保留批次维度。如果 False,使用选择而不是切片。此功能应与 exclusive=True 一起使用。

  • workers (sequence of int or str, or None, optional) – 用于将批次分区的工作者。 如果 None,批次将被分区到世界中的所有工作者。 (默认值: None)

Return type:

Callable[..., Sequence[Tuple[int, Optional[Tuple[Any, ...]], Optional[Dict[str, Any]]]]]

Returns:

一个分区函数。

torchopt.distributed.batch_partitioner(*args: Any, **kwargs: Any) list[tuple[int, Args | None, KwArgs | None]]

用于批处理维度的分区器。将参数沿第一个维度分割并复制到所有工作节点。

批次将被划分为min(batch_size, num_workers)个分区,其中 num_workers是世界中的工作线程数。 当batch_size > num_workers时,可以在单个RPC调用中前向传递多个批次样本。

argskwargs中的所有张量将沿着维度dim进行分区,而非张量值将被广播到各个分区。

torchopt.distributed.mean_reducer(results)[source]

通过平均来减少结果。

Return type:

Tensor

torchopt.distributed.sum_reducer(results)[source]

通过求和来减少结果。

Return type:

Tensor

函数并行化包装器

parallelize([partitioner, reducer, timeout])

返回一个用于并行化函数的装饰器。

parallelize_async([partitioner, reducer, ...])

返回一个用于并行化函数的装饰器。

parallelize_sync([partitioner, reducer, timeout])

返回一个用于并行化函数的装饰器。

torchopt.distributed.parallelize(partitioner=None, reducer=None, timeout=-1.0)[source]

返回一个用于并行化函数的装饰器。

此装饰器可用于在多个工作线程之间并行化函数调用。

Parameters:
  • partitioner (int, str, 或 callable, 可选) – 一个将参数分区到多个工作节点的分区器。(默认: batch_partitioner())

  • reducer (callableNone, 可选) – 一个用于减少多个工作者结果的reducer。如果 None,则不减少结果。(默认值:None

  • timeout (float, 可选) – RPC调用的超时时间。 (默认: rpc.api.UNSET_RPC_TIMEOUT)

Return type:

Callable[[Callable[..., TypeVar(T)]], Callable[..., Union[list[TypeVar(T)], TypeVar(U)]]]

Returns:

装饰器函数。

torchopt.distributed.parallelize_async(partitioner=None, reducer=None, timeout=-1.0)[source]

返回一个用于并行化函数的装饰器。

此装饰器可用于在多个工作节点上并行化函数调用。该函数将在远程工作节点上异步调用。装饰后的函数将返回结果的torch.Future实例。

Parameters:
  • partitioner (int, str, 或 callable, 可选) – 一个将参数分区到多个工作节点的分区器。(默认: batch_partitioner())

  • reducer (callableNone, 可选) – 一个用于减少多个工作者结果的reducer。如果 None,则不减少结果。(默认值:None

  • timeout (float, 可选) – RPC调用的超时时间。 (默认: rpc.api.UNSET_RPC_TIMEOUT)

Return type:

Callable[[Callable[..., TypeVar(T)]], Callable[..., Union[Future[list[TypeVar(T)]], Future[TypeVar(U)]]]]

Returns:

装饰器函数。

torchopt.distributed.parallelize_sync(partitioner=None, reducer=None, timeout=-1.0)

返回一个用于并行化函数的装饰器。

此装饰器可用于在多个工作线程之间并行化函数调用。

Parameters:
  • partitioner (int, str, 或 callable, 可选) – 一个将参数分区到多个工作节点的分区器。(默认: batch_partitioner())

  • reducer (callableNone, 可选) – 一个用于减少多个工作者结果的reducer。如果 None,则不减少结果。(默认值:None

  • timeout (float, 可选) – RPC调用的超时时间。 (默认: rpc.api.UNSET_RPC_TIMEOUT)

Return type:

Callable[[Callable[..., TypeVar(T)]], Callable[..., Union[list[TypeVar(T)], TypeVar(U)]]]

Returns:

装饰器函数。

分布式自动求导

context()

在使用分布式自动求导时,用于包装前向和后向传递的上下文对象。

get_gradients(context_id)

检索从Tensor到该Tensor的适当梯度的映射,该梯度在提供的上下文中累积,对应于给定的context_id,作为分布式自动梯度反向传递的一部分。

backward(autograd_ctx_id, tensors[, ...])

为本地参数执行分布式反向传播。

grad(autograd_ctx_id, outputs, inputs[, ...])

计算并返回输出相对于输入的梯度之和。

torchopt.distributed.autograd.context()[source]

上下文对象,用于在使用分布式自动梯度时包装前向和后向传递。在with语句中生成的context_id是必需的,以在所有工作节点上唯一标识一个分布式后向传递。每个工作节点存储与此context_id相关的元数据,这些元数据是正确执行分布式自动梯度传递所必需的。

Example::
>>> # xdoctest: +SKIP
>>> import torch.distributed.autograd as dist_autograd
>>> with dist_autograd.context() as context_id:
>>>     t1 = torch.rand((3, 3), requires_grad=True)
>>>     t2 = torch.rand((3, 3), requires_grad=True)
>>>     loss = rpc.rpc_sync("worker1", torch.add, args=(t1, t2)).sum()
>>>     dist_autograd.backward(context_id, [loss])
torchopt.distributed.autograd.get_gradients(context_id: int) Dict[Tensor, Tensor]

从Tensor中检索一个映射,该映射对应于给定context_id的适当梯度,该梯度在提供的上下文中累积,作为分布式自动梯度反向传递的一部分。

Parameters:

context_id (int) – 我们应该为其检索梯度的自动求导上下文ID。

Returns:

一个映射,其中键是Tensor,值是该Tensor的关联梯度。

Example::
>>> import torch.distributed.autograd as dist_autograd
>>> with dist_autograd.context() as context_id:
>>>     t1 = torch.rand((3, 3), requires_grad=True)
>>>     t2 = torch.rand((3, 3), requires_grad=True)
>>>     loss = t1 + t2
>>>     dist_autograd.backward(context_id, [loss.sum()])
>>>     grads = dist_autograd.get_gradients(context_id)
>>>     print(grads[t1])
>>>     print(grads[t2])
torchopt.distributed.autograd.backward(autograd_ctx_id, tensors, retain_graph=False, inputs=None)[source]

为本地参数执行分布式反向传递。

计算给定张量相对于图叶节点的梯度之和。

Parameters:
  • autograd_ctx_id (int) – 自动求导上下文ID。

  • 张量 (Tensor序列Tensor) – 将计算其导数的张量。

  • retain_graph (bool, 可选) – 如果为False,用于计算梯度的图将被释放。请注意,在几乎所有情况下,将此选项设置为True是不必要的,通常可以通过更高效的方式解决。 (默认: False)

  • inputs (Tensor, sequence of Tensor, or None, optional) – 输入相对于这些输入的梯度将被累积到 .grad 中。所有其他张量将被忽略。如果未提供,梯度将累积到用于计算 tensors 的所有叶张量中。(默认值:None

Return type:

torchopt.distributed.autograd.grad(autograd_ctx_id, outputs, inputs, retain_graph=False, allow_unused=False)[source]

计算并返回输出相对于输入的梯度之和。

Parameters:
  • autograd_ctx_id (int) – 自动求导上下文ID。

  • 输出 (张量张量序列) – 微分函数的输出。

  • inputs (TensorTensor 的序列) – 输入,关于这些输入的梯度将被返回 (并且不会累积到 .grad 中)。

  • retain_graph (bool, 可选) – 如果为False,用于计算梯度的图将被释放。请注意,在几乎所有情况下,将此选项设置为True是不必要的,并且通常可以通过更高效的方式解决。 (默认值: False)

  • allow_unused (bool, 可选) – 如果为False,指定未在计算输出时使用的输入(因此它们的梯度始终为零)将导致错误。 (默认值: False)

Return type:

Tuple[Optional[Tensor], ...]

通用工具

extract_state_dict(target, *[, by, device, ...])

提取目标状态。

recover_state_dict(target, state)

恢复状态。

stop_gradient(target)

停止输入对象的梯度。

提取状态字典

torchopt.extract_state_dict(target, *, by='reference', device=None, with_buffers=True, detach_buffers=False, enable_visual=False, visual_prefix='')[source]

提取目标状态。

由于张量使用grad_fn将其自身与前一个计算图连接,反向传播的梯度将流过该张量,并继续流向由grad_fn连接的张量。某些算法需要手动将张量从计算图中分离。

请注意,提取的状态是一个引用,这意味着任何就地操作符都会影响从中提取状态的目标。

Parameters:
  • 目标 (nn.ModuleMetaOptimizer) – 它可以是一个 nn.Moduletorchopt.MetaOptimizer

  • by (str, 可选) – 目标中张量的提取策略。(默认:'reference') - 'reference': 提取的张量将是原始张量的引用。 - 'copy': 提取的张量将是原始张量的克隆。这使得复制的张量具有指向原始张量的grad_fn,即函数。 - 'deepcopy': 提取的张量将从原始张量进行深拷贝。深拷贝的张量将脱离原始计算图。

  • device (DeviceNone, 可选) – 如果指定,将提取的状态移动到指定的设备。(默认: None)

  • with_buffers (bool, 可选) – 提取缓冲区与参数一起,此参数仅在输入目标为 nn.Module 时使用。(默认值:True

  • detach_buffers (bool, optional) – 是否分离对缓冲区的引用,此参数仅在输入目标为nn.Moduleby='reference'时使用。(默认值:False

  • enable_visual (bool, 可选) – 添加额外的注释,这些注释可以用于计算图的可视化。目前,这个标志只对nn.Module有效,但我们稍后将支持torchopt.MetaOptimizer。(默认值:False

  • visual_prefix (str, optional) – 可视化注释的前缀。 (default: '')

Return type:

ModuleState | tuple[Union[Tensor, Tuple[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]], ...], List[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Dict[Any, Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], Deque[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]], CustomTreeNode[Union[Tensor, Tuple[TensorTree, ...], List[TensorTree], Dict[Any, TensorTree], Deque[TensorTree], CustomTreeNode[TensorTree]]]], ...]

Returns:

从输入对象中提取的状态。

恢复状态字典

torchopt.recover_state_dict(target, state)[source]

恢复状态。

此函数与extract_state兼容。

请注意,恢复过程不是原地进行的,因此对象的张量不会被修改。

Parameters:
  • target (nn.ModuleMetaOptimizer) – 需要恢复的目标。

  • state (ModuleState序列Tensor) – 恢复状态。

Return type:

停止梯度

torchopt.stop_gradient(target)[source]

停止输入对象的梯度。

由于张量使用grad_fn将其自身与前一个计算图连接,反向传播的梯度将通过张量流动,并继续流向由grad_fn连接的张量。某些算法需要手动将张量从计算图中分离。

请注意,stop_gradient() 操作是原地进行的。

Parameters:

target (ModuleState, nn.Module, MetaOptimizer, 或 treeTensor) – 要从计算图中分离的目标,它可以是 nn.Module, torchopt.MetaOptimizer, torchopt.MetaOptimizer 的状态,或者只是一个普通的张量列表。

Return type:

可视化梯度流

make_dot(var[, params, show_attrs, ...])

生成PyTorch自动微分图的Graphviz表示。

制作点

torchopt.visual.make_dot(var, params=None, show_attrs=False, show_saved=False, max_attr_chars=50)[source]

生成PyTorch autograd图的Graphviz表示。

如果一个节点代表一个反向函数,它是灰色的。否则,节点代表一个张量,并且是蓝色、橙色或绿色:

  • Blue

    需要梯度的可达叶子张量(在backward()期间将填充其grad字段的张量)。

  • Orange

    保存自定义自动梯度函数的张量以及由内置反向节点保存的张量。

  • Green

    作为输出传递的张量。

  • Dark green

    如果任何输出是一个视图,我们用深绿色节点表示其基础张量。

Parameters:
  • var (Tensor序列Tensor) – 输出张量。

  • params (Optional[Mapping[str, Tensor], ModuleState, Generator, Iterable[Union[Mapping[str, Tensor], ModuleState, Generator]]]]) – (dict[str, Tensor], ModuleState, iterable of tuple[str, Tensor], or None, optional): 用于为需要梯度的节点添加名称的参数。(默认值: None)

  • show_attrs (bool, 可选) – 是否显示反向节点的非张量属性。 (默认: False)

  • show_saved (bool, 可选) – 是否显示非自定义自动求导函数保存的张量节点。自定义函数保存的张量节点(如果存在)将始终显示。(默认值:False

  • max_attr_chars (int, 可选) – 如果 show_attrsTrue,设置显示任何给定属性的最大字符数。(默认值:50

Return type:

有向图