优化器

TorchOpt的核心设计遵循函数式编程的理念。与functorch一致，用户可以在PyTorch中使用模型、优化器和训练进行函数式编程。我们首先介绍我们的函数式优化器，它将优化过程视为一种函数转换。

功能优化器

目前，TorchOpt 支持 4 种功能优化器：sgd()、adam()、rmsprop() 和 adamw()。

`torchopt.FuncOptimizer`(impl, *[, inplace])	一个包装类，用于持有功能优化器。
`torchopt.adadelta`([lr, rho, eps, ...])	创建一个AdaDelta优化器的功能版本。
`torchopt.adagrad`([lr, lr_decay, ...])	创建一个AdaGrad优化器的功能版本。
`torchopt.adam`([lr, betas, eps, ...])	创建一个Adam优化器的功能版本。
`torchopt.adamw`([lr, betas, eps, ...])	创建一个带有权重衰减正则化的Adam优化器的功能版本。
`torchopt.adamax`([lr, betas, eps, ...])	创建一个AdaMax优化器的功能版本。
`torchopt.radam`([lr, betas, eps, ...])	创建一个RAdam优化器的功能版本。
`torchopt.rmsprop`([lr, alpha, eps, ...])	创建一个RMSProp优化器的功能版本。
`torchopt.sgd`(lr[, momentum, dampening, ...])	创建一个规范随机梯度下降优化器的功能版本。

应用参数更新

TorchOpt 通过将梯度和优化器状态传递给优化器函数来应用更新，提供功能性 API。

torchopt.apply_updates(params, updates, *[, ...])

将更新应用于相应的参数。

这是一个结合了functorch的功能优化示例：

class Net(nn.Module): ...

class Loader(DataLoader): ...

net = Net()  # init
loader = Loader()
optimizer = torchopt.adam(lr)

model, params = functorch.make_functional(net)           # use functorch extract network parameters
opt_state = optimizer.init(params)                       # init optimizer

xs, ys = next(loader)                                    # get data
pred = model(params, xs)                                 # forward
loss = F.cross_entropy(pred, ys)                         # compute loss

grads = torch.autograd.grad(loss, params)                # compute gradients
updates, opt_state = optimizer.update(grads, opt_state)  # get updates
params = torchopt.apply_updates(params, updates)         # update network parameters

我们还提供了一个包装器 torchopt.FuncOptimizer 以使维护优化器状态更容易：

net = Net()  # init
loader = Loader()
optimizer = torchopt.FuncOptimizer(torchopt.adam())      # wrap with `torchopt.FuncOptimizer`

model, params = functorch.make_functional(net)           # use functorch extract network parameters

for xs, ys in loader:                                    # get data
    pred = model(params, xs)                             # forward
    loss = F.cross_entropy(pred, ys)                     # compute loss

    params = optimizer.step(loss, params)                # update network parameters

经典面向对象编程优化器

结合上述功能优化器，我们可以定义经典的面向对象编程（OOP）优化器。我们设计了一个基类 torchopt.Optimizer，其接口与 torch.optim.Optimizer 相同。我们提供了原始的 PyTorch API（例如，zero_grad() 或 step()）用于传统的类似 PyTorch（OOP）的参数更新。

`torchopt.Optimizer`(params, impl)	一个类似于 `torch.optim.Optimizer` 的经典优化器基类。
`torchopt.AdaDelta`(params[, lr, rho, eps, ...])	经典的AdaDelta优化器。
`torchopt.Adadelta`	`AdaDelta` 的别名
`torchopt.AdaGrad`(params[, lr, lr_decay, ...])	经典的AdaGrad优化器。
`torchopt.Adagrad`	`AdaGrad` 的别名
`torchopt.Adam`(params[, lr, betas, eps, ...])	经典的Adam优化器。
`torchopt.AdamW`(params[, lr, betas, eps, ...])	经典的AdamW优化器。
`torchopt.AdaMax`(params[, lr, betas, eps, ...])	经典的AdaMax优化器。
`torchopt.Adamax`	`AdaMax` 的别名
`torchopt.RAdam`(params[, lr, betas, eps, ...])	经典的RAdam优化器。
`torchopt.RMSProp`(params[, lr, alpha, eps, ...])	经典的RMSProp优化器。
`torchopt.SGD`(params, lr[, momentum, ...])	经典的SGD优化器。

通过将低级API torchopt.Optimizer 与之前的功能优化器结合，我们可以实现高级API：

learning_rate = 1.0
# High-level API
optim = torchopt.Adam(net.parameters(), lr=learning_rate)
# which can be achieved by low-level API:
optim = torchopt.Optimizer(net.parameters(), torchopt.adam(lr=learning_rate))

以下是一个类似PyTorch的API示例：

net = Net()  # init
loader = Loader()
optimizer = torchopt.Adam(net.parameters())

xs, ys = next(loader)             # get data
pred = net(xs)                    # forward
loss = F.cross_entropy(pred, ys)  # compute loss

optimizer.zero_grad()             # zero gradients
loss.backward()                   # backward
optimizer.step()                  # step updates

组合转换

用户在进行最终更新之前，总是需要进行多次梯度变换（函数）。在TorchOpt的设计中，我们将这些函数视为torchopt.chain()的派生。因此，我们可以构建自己的链，如torchopt.chain(torchopt.clip_grad_norm(max_norm=1.), torchopt.sgd(lr=1., moment_requires_grad=True))，以裁剪梯度并使用sgd()更新参数。

torchopt.chain(*transformations)

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

注意

torchopt.chain() 将依次进行转换，因此顺序很重要。例如，我们需要先进行梯度归一化，然后进行优化器步骤。在 torchopt.chain() 函数中，顺序应为 (clip, sgd)。

这里是一个将torchopt.clip_grad_norm()和torchopt.adam()链接在一起的示例，用于功能优化器和面向对象优化器。

func_optimizer = torchopt.chain(torchopt.clip_grad_norm(max_norm=2.0), torchopt.adam(1e-1))
oop_optimizer = torchopt.Optimizer(net.parameters() func_optimizer)

优化器钩子

用户还可以添加优化器钩子来控制梯度流。

`torchopt.hook.register_hook`(hook)	无状态的身份转换，保持输入梯度不变。
`torchopt.hook.zero_nan_hook`(g)	将 `nan` 替换为零。
`torchopt.hook.nan_to_num_hook`([nan, posinf, ...])	返回一个`nan`到num的钩子，用于将`nan` / `+inf` / `-inf`替换为给定的数字。

例如，torchopt.hook.zero_nan_hook() 注册钩子到一阶梯度。在反向传播过程中，NaN 梯度将被设置为 0。这里是一个与 torchopt.chain() 结合使用的操作示例。

impl = torchopt.chain(torchopt.hook.register_hook(torchopt.hook.zero_nan_hook), torchopt.adam(1e-1))

优化器调度

TorchOpt 还提供了学习率调度器的实现，这些调度器可用于在训练过程中控制学习率。 TorchOpt 主要提供线性学习率调度器和多项式学习率调度器。

`torchopt.schedule.linear_schedule`(...[, ...])	为了方便起见，将多项式调度别名为线性调度。
`torchopt.schedule.polynomial_schedule`(...[, ...])	构建一个从初始值到结束值具有多项式过渡的调度。

这是一个将优化器与学习率调度器结合的例子。

functional_adam = torchopt.adam(
    lr=torchopt.schedule.linear_schedule(
        init_value=1e-3, end_value=1e-4, transition_steps=10000, transition_begin=2000
    )
)

adam = torchopt.Adam(
    net.parameters(),
    lr=torchopt.schedule.linear_schedule(
        init_value=1e-3, end_value=1e-4, transition_steps=10000, transition_begin=2000
    ),
)

笔记本教程

查看笔记本教程在Functional Optimizer。