torch.optim¶

torch.optim 是一个实现各种优化算法的包。

最常用的方法已经得到支持，并且接口足够通用，因此将来也可以轻松集成更复杂的方法。

如何使用优化器¶

要使用 torch.optim，您必须构建一个优化器对象，该对象将保存当前状态并根据计算出的梯度更新参数。

构建它¶

要构造一个Optimizer，你必须给它一个包含参数（所有都应该是Variable）的可迭代对象来优化。然后，你可以指定优化器特定的选项，例如学习率、权重衰减等。

示例：

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = optim.Adam([var1, var2], lr=0.0001)

每个参数的选项¶

Optimizer 还支持指定每个参数的选项。为此，不要传递 Variable 的可迭代对象，而是传递一个 dict 的可迭代对象。每个字典将定义一个单独的参数组，并且应该包含一个 params 键，其中包含属于该组的参数列表。其他键应与优化器接受的参数关键字匹配，并将用作该组的优化选项。

例如，当需要为每一层指定学习率时，这非常有用：

optim.SGD([
                {'params': model.base.parameters(), 'lr': 1e-2},
                {'params': model.classifier.parameters()}
            ], lr=1e-3, momentum=0.9)

这意味着 model.base 的参数将使用 1e-2 的学习率，而 model.classifier 的参数将保持默认的学习率 1e-3。最后，所有参数将使用 0.9 的动量。

注意

您仍然可以将选项作为关键字参数传递。它们将被用作默认值，在未覆盖它们的组中使用。当您只想改变一个选项，同时保持所有其他选项在参数组之间一致时，这非常有用。

还请考虑以下与参数不同惩罚相关的示例。请记住，parameters() 返回一个包含所有可学习参数的可迭代对象，包括偏差和其他可能需要不同惩罚的参数。为了解决这个问题，可以为每个参数组指定单独的惩罚权重：

bias_params = [p for name, p in self.named_parameters() if 'bias' in name]
others = [p for name, p in self.named_parameters() if 'bias' not in name]

optim.SGD([
                {'params': others},
                {'params': bias_params, 'weight_decay': 0}
            ], weight_decay=1e-2, lr=1e-2)

通过这种方式，偏置项与非偏置项被隔离，并且为偏置项特别设置了weight_decay 为0，以避免对该组进行任何惩罚。

执行优化步骤¶

所有优化器都实现了一个step()方法，用于更新参数。它可以通过两种方式使用：

`optimizer.step()`¶

这是大多数优化器支持的简化版本。该函数可以在使用例如backward()计算梯度后调用。

示例：

for input, target in dataset:
    optimizer.zero_grad()
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()

`optimizer.step(closure)`¶

一些优化算法，如共轭梯度和LBFGS，需要多次重新评估函数，因此你必须传入一个闭包，允许它们重新计算你的模型。闭包应该清除梯度，计算损失，并返回它。

示例：

for input, target in dataset:
    def closure():
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)
        loss.backward()
        return loss
    optimizer.step(closure)

基类¶

class torch.optim.Optimizer(params, defaults)[源代码]¶

所有优化器的基类。

警告

参数需要指定为具有确定性顺序的集合，该顺序在运行之间保持一致。不满足这些属性的对象示例包括集合和对字典值的迭代器。

Parameters

参数 (可迭代对象) – 一个包含 torch.Tensor 或 dict 的可迭代对象。指定应优化的张量。
默认值 (字典[字符串, 任意]) – (字典): 包含优化选项默认值的字典（当参数组未指定这些值时使用）。

`Optimizer.add_param_group`	向 `Optimizer` 的 param_groups 添加一个参数组。
`Optimizer.load_state_dict`	加载优化器状态。
`Optimizer.state_dict`	返回优化器的状态为一个`dict`。
`Optimizer.step`	执行单个优化步骤（参数更新）。
`Optimizer.zero_grad`	重置所有优化的 `torch.Tensor` 的梯度。

算法¶

`Adadelta`	实现Adadelta算法。
`Adagrad`	实现Adagrad算法。
`Adam`	实现Adam算法。
`AdamW`	实现 AdamW 算法。
`SparseAdam`	SparseAdam 实现了适用于稀疏梯度的 Adam 算法的掩码版本。
`Adamax`	实现Adamax算法（基于无穷范数的Adam变种）。
`ASGD`	实现平均随机梯度下降。
`LBFGS`	实现L-BFGS算法。
`NAdam`	实现NAdam算法。
`RAdam`	实现RAdam算法。
`RMSprop`	实现RMSprop算法。
`Rprop`	实现弹性反向传播算法。
`SGD`	实现随机梯度下降（可选地带有动量）。

我们的许多算法都有多种实现，针对性能、可读性和/或通用性进行了优化，因此如果没有特别指定实现方式，我们会尝试默认选择当前设备上最快的一般实现。

我们有3种主要的实现类别：for-loop、foreach（多张量）和fused。最直接的实现方式是使用for-loop遍历参数并进行大量计算。for-loop通常比我们的foreach实现要慢，foreach实现将参数组合成一个多张量，并一次性运行大量计算，从而节省了许多顺序内核调用。我们的一些优化器甚至有更快的fused实现，它们将大量计算融合到一个内核中。我们可以将foreach实现视为水平融合，而fused实现则是在此基础上的垂直融合。

一般来说，这3种实现的性能排序是 fused > foreach > for-loop。因此，在适用的情况下，我们默认使用foreach而不是for-loop。适用意味着foreach实现是可用的，用户没有指定任何特定于实现的kwargs（例如，fused、foreach、differentiable），并且所有张量都是本地的并且位于CUDA上。请注意，虽然fused应该比foreach更快，但这些实现是较新的，我们希望在全面切换之前给它们更多的时间进行测试。欢迎您尝试使用它们！

以下是一个表格，展示了每种算法的可用和默认实现：

算法	默认	有 foreach 吗？	已融合？
`Adadelta`	foreach	是	否
`Adagrad`	遍历	是	否
`Adam`	遍历	是	是
`AdamW`	遍历	是	是
`SparseAdam`	for循环	否	否
`Adamax`	遍历	是	否
`ASGD`	foreach	是	否
`LBFGS`	for循环	否	否
`NAdam`	foreach	是	否
`RAdam`	foreach	是	否
`RMSprop`	foreach	是	否
`Rprop`	foreach	是	否
`SGD`	foreach	是	否

如何调整学习率¶

torch.optim.lr_scheduler 提供了几种根据epoch数量调整学习率的方法。torch.optim.lr_scheduler.ReduceLROnPlateau 允许基于某些验证指标动态减少学习率。

学习率调度应在优化器的更新之后应用；例如，你应该这样编写代码：

示例：

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
scheduler = ExponentialLR(optimizer, gamma=0.9)

for epoch in range(20):
    for input, target in dataset:
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()
    scheduler.step()

大多数学习率调度器可以连续调用（也称为链式调度器）。结果是每个调度器依次应用于前一个调度器获得的学习率。

示例：

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
scheduler1 = ExponentialLR(optimizer, gamma=0.9)
scheduler2 = MultiStepLR(optimizer, milestones=[30,80], gamma=0.1)

for epoch in range(20):
    for input, target in dataset:
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()
    scheduler1.step()
    scheduler2.step()

在文档的许多地方，我们将使用以下模板来引用调度器算法。

>>> scheduler = ...
>>> for epoch in range(100):
>>>     train(...)
>>>     validate(...)
>>>     scheduler.step()

警告

在 PyTorch 1.1.0 之前，学习率调度器需要在优化器的更新之前调用；1.1.0 以一种不兼容的方式改变了这一行为。如果你在使用学习率调度器（调用 scheduler.step()）之前调用优化器的更新（调用 optimizer.step()），这将跳过学习率调度的第一个值。如果你在升级到 PyTorch 1.1.0 后无法重现结果，请检查你是否在错误的时间调用了 scheduler.step()。

`lr_scheduler.LambdaLR`	将每个参数组的学习率设置为初始学习率乘以给定函数。
`lr_scheduler.MultiplicativeLR`	将每个参数组的学习率乘以指定函数中给出的因子。
`lr_scheduler.StepLR`	每个参数组的学习率每经过step_size个epochs衰减gamma倍。
`lr_scheduler.MultiStepLR`	当训练轮数达到某个里程碑时，每个参数组的学习率会按gamma衰减。
`lr_scheduler.ConstantLR`	将每个参数组的学习率乘以一个小的常数因子，直到训练轮数达到预定义的里程碑：total_iters。
`lr_scheduler.LinearLR`	每个参数组的学习率通过线性改变一个小乘法因子来衰减，直到epoch数达到预定义的里程碑：total_iters。
`lr_scheduler.ExponentialLR`	每个参数组的学习率每轮衰减gamma。
`lr_scheduler.PolynomialLR`	使用给定的 total_iters 中的多项式函数衰减每个参数组的学习率。
`lr_scheduler.CosineAnnealingLR`	使用余弦退火计划为每个参数组设置学习率，其中 $\eta_{max}$ 设置为初始学习率， $T_{cur}$ 是从上次重启以来的轮数，在 SGDR 中：
`lr_scheduler.ChainedScheduler`	链式学习率调度器列表。
`lr_scheduler.SequentialLR`	接收在优化过程中预期按顺序调用的调度器列表，以及提供精确间隔的里程碑点，以反映在给定时期应调用哪个调度器。
`lr_scheduler.ReduceLROnPlateau`	当某个指标停止改进时，降低学习率。
`lr_scheduler.CyclicLR`	根据循环学习率策略（CLR）设置每个参数组的学习率。
`lr_scheduler.OneCycleLR`	根据1cycle学习率策略设置每个参数组的学习率。
`lr_scheduler.CosineAnnealingWarmRestarts`	使用余弦退火计划为每个参数组设置学习率，其中 $\eta_{max}$ 设置为初始学习率， $T_{cur}$ 是从上次重启以来的轮数， $T_{i}$ 是在 SGDR 中两次热重启之间的轮数：

权重平均（SWA 和 EMA）¶

torch.optim.swa_utils 实现了随机权重平均（SWA）和指数移动平均（EMA）。特别是， torch.optim.swa_utils.AveragedModel 类实现了SWA和EMA模型， torch.optim.swa_utils.SWALR 实现了SWA学习率调度器， torch.optim.swa_utils.update_bn() 是一个用于在训练结束时更新SWA/EMA批量归一化统计信息的实用函数。

SWA 已在平均权重导致更宽的最优解和更好的泛化中被提出。

EMA 是一种广泛使用的技术，通过减少所需的权重更新次数来缩短训练时间。它是 Polyak 平均的一种变体，但在迭代过程中使用指数权重而不是相等的权重。

构建平均模型¶

类用于计算SWA或EMA模型的权重。

您可以通过运行以下命令来创建一个SWA平均模型：

>>> averaged_model = AveragedModel(model)

EMA 模型通过将 multi_avg_fn 参数指定如下进行构建：

>>> decay = 0.999
>>> averaged_model = AveragedModel(model, multi_avg_fn=get_ema_multi_avg_fn(decay))

Decay 是一个介于 0 和 1 之间的参数，用于控制平均参数的衰减速度。如果没有提供给 get_ema_multi_avg_fn，默认值为 0.999。

get_ema_multi_avg_fn 返回一个函数，该函数将以下EMA方程应用于权重：

W^\textrm{EMA}_{t+1} = \alpha W^\textrm{EMA}_{t} + (1 - \alpha) W^\textrm{model}_t

其中 alpha 是 EMA 衰减。

这里的模型 model 可以是任意的 torch.nn.Module 对象。averaged_model 将跟踪 model 参数的运行平均值。要更新这些平均值，您应该在 optimizer.step() 之后使用 update_parameters() 函数：

>>> averaged_model.update_parameters(model)

对于SWA和EMA，这个调用通常在优化器step()之后立即执行。在SWA的情况下，通常会在训练开始时跳过一些步骤。

自定义平均策略¶

默认情况下，torch.optim.swa_utils.AveragedModel 计算您提供的参数的运行平均值，但您也可以使用自定义平均函数，通过 avg_fn 或 multi_avg_fn 参数：

avg_fn 允许定义一个函数，该函数对每个参数元组（平均参数，模型参数）进行操作，并应返回新的平均参数。
multi_avg_fn 允许定义更高效的操作，这些操作同时作用于参数列表的元组（平均参数列表，模型参数列表），例如使用 torch._foreach* 函数。此函数必须就地更新平均参数。

在下面的示例中，ema_model 使用 avg_fn 参数计算指数移动平均值：

>>> ema_avg = lambda averaged_model_parameter, model_parameter, num_averaged:\
>>>         0.9 * averaged_model_parameter + 0.1 * model_parameter
>>> ema_model = torch.optim.swa_utils.AveragedModel(model, avg_fn=ema_avg)

在下面的示例中，ema_model 使用更高效的 multi_avg_fn 参数计算指数移动平均值：

>>> ema_model = AveragedModel(model, multi_avg_fn=get_ema_multi_avg_fn(0.9))

SWA 学习率调度¶

通常，在SWA中，学习率被设置为一个较高的常数值。SWALR 是一个学习率调度器，它将学习率退火到一个固定值，然后保持不变。例如，以下代码创建了一个调度器，该调度器在每个参数组中，在5个周期内线性地将学习率从其初始值退火到0.05：

>>> swa_scheduler = torch.optim.swa_utils.SWALR(optimizer, \
>>>         anneal_strategy="linear", anneal_epochs=5, swa_lr=0.05)

您也可以通过设置anneal_strategy="cos"来使用余弦退火到固定值，而不是线性退火。

处理批量归一化¶

update_bn() 是一个实用函数，允许在训练结束时为 SWA 模型在给定的数据加载器 loader 上计算批归一化统计数据：

>>> torch.optim.swa_utils.update_bn(loader, swa_model)

update_bn() 将 swa_model 应用于数据加载器中的每个元素，并计算模型中每个批量归一化层的激活统计数据。

警告

update_bn() 假设数据加载器 loader 中的每个批次要么是一个张量，要么是一个张量列表，其中第一个元素是网络 swa_model 应该应用的张量。如果你的数据加载器具有不同的结构，你可以通过在数据集的每个元素上使用 swa_model 进行前向传递来更新 swa_model 的批量归一化统计数据。

综合运用：SWA¶

在下面的示例中，swa_model 是累积权重平均值的SWA模型。我们总共训练模型300个周期，并在第160个周期切换到SWA学习率计划，并开始收集参数的SWA平均值：

>>> loader, optimizer, model, loss_fn = ...
>>> swa_model = torch.optim.swa_utils.AveragedModel(model)
>>> scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=300)
>>> swa_start = 160
>>> swa_scheduler = SWALR(optimizer, swa_lr=0.05)
>>>
>>> for epoch in range(300):
>>>       for input, target in loader:
>>>           optimizer.zero_grad()
>>>           loss_fn(model(input), target).backward()
>>>           optimizer.step()
>>>       if epoch > swa_start:
>>>           swa_model.update_parameters(model)
>>>           swa_scheduler.step()
>>>       else:
>>>           scheduler.step()
>>>
>>> # 在最后更新swa_model的bn统计信息
>>> torch.optim.swa_utils.update_bn(loader, swa_model)
>>> # 使用swa_model对测试数据进行预测
>>> preds = swa_model(test_input)

综合运用：EMA¶

在下面的示例中，ema_model 是用于累积权重指数衰减平均值的EMA模型，衰减率为0.999。我们总共训练模型300个epoch，并立即开始收集EMA平均值。

>>> loader, optimizer, model, loss_fn = ...
>>> ema_model = torch.optim.swa_utils.AveragedModel(model, \
>>>             multi_avg_fn=torch.optim.swa_utils.get_ema_multi_avg_fn(0.999))
>>>
>>> for epoch in range(300):
>>>       for input, target in loader:
>>>           optimizer.zero_grad()
>>>           loss_fn(model(input), target).backward()
>>>           optimizer.step()
>>>           ema_model.update_parameters(model)
>>>
>>> # 在最后更新ema_model的bn统计信息
>>> torch.optim.swa_utils.update_bn(loader, ema_model)
>>> # 使用ema_model对测试数据进行预测
>>> preds = ema_model(test_input)