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 的可迭代对象。每个 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)

所有优化器的基类。

警告

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

参数

• params (iterable) – torch.Tensor 或 dict 的可迭代对象。指定应优化哪些张量。

• defaults (Dict[str, Any]) – (dict)：包含优化选项的默认值（在参数组未指定这些值时使用）的字典。

Optimizer.add_param_group	向 Optimizer 的 param_groups 添加一个参数组。
Optimizer.load_state_dict	加载优化器状态。
Optimizer.state_dict	将优化器的状态作为 dict 返回。
Optimizer.step	执行单个优化步骤（参数更新）。
优化器.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 循环、foreach（多张量）和融合。最直接的实现是对具有大量计算块的参数进行 for 循环。For 循环通常比我们的 foreach 实现慢，后者将参数组合成多张量并一次运行大量计算块，从而节省了许多顺序内核调用。我们的几个优化器甚至有更快的融合实现，它将大量计算块融合到一个内核中。我们可以将 foreach 实现视为水平融合，并将融合实现视为在其之上垂直融合。

通常，3 种实现的性能顺序为融合 > foreach > for 循环。因此，在适用时，我们默认为 foreach over for 循环。适用意味着 foreach 实现可用，用户未指定任何特定于实现的 kwargs（例如，融合、foreach、可微），并且所有张量都是本机张量且在 CUDA 上。请注意，虽然融合应该比 foreach 还要快，但这些实现较新，我们希望在到处翻转开关之前给它们更多的时间进行内部构建。不过，欢迎您试用它们！

下表显示了每种算法的可用实现和默认实现

算法	默认	有 foreach？	有 fused？
Adadelta	foreach	是	否
Adagrad	foreach	是	否
Adam	foreach	是	是
AdamW	foreach	是	是
SparseAdam	for-loop	否	否
Adamax	foreach	是	否
ASGD	foreach	是	否
LBFGS	for-loop	否	否
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 以破坏向后兼容性的方式更改了此行为。如果您在优化器的更新（调用 optimizer.step()）之前使用学习率计划器（调用 scheduler.step()），这将跳过学习率计划的第一个值。如果您在升级到 PyTorch 1.1.0 后无法重现结果，请检查您是否在错误的时间调用了 scheduler.step()。

lr_scheduler.LambdaLR	将每个参数组的学习率设置为初始 lr 乘以给定函数。
lr_scheduler.MultiplicativeLR	将每个参数组的学习率乘以指定函数中给出的因子。
lr_scheduler.StepLR	每隔 step_size 个 epoch 将每个参数组的学习率衰减 gamma。
lr_scheduler.MultiStepLR	一旦 epoch 数量达到某个里程碑，就将每个参数组的学习率衰减 gamma。
lr_scheduler.ConstantLR	将每个参数组的学习率乘以一个小的常数因子，直到 epoch 数达到预定义的里程碑：total_iters。
lr_scheduler.LinearLR	通过线性改变小的乘法因子来衰减每个参数组的学习率，直到 epoch 数达到预定义的里程碑：total_iters。
lr_scheduler.ExponentialLR	每个 epoch 将每个参数组的学习率衰减 gamma。
lr_scheduler.PolynomialLR	使用给定的 total_iters 中的多项式函数衰减每个参数组的学习率。
lr_scheduler.CosineAnnealingLR	使用余弦退火计划设置每个参数组的学习率，其中 $\eta_{max}$ 设置为初始 lr，$T_{cur}$ 是 SGDR 中自上次重新启动以来的 epoch 数
lr_scheduler.ChainedScheduler	链接学习率计划程序列表。
lr_scheduler.SequentialLR	接收在优化过程中按顺序调用的计划程序列表和里程碑点，这些点提供准确的间隔来反映在给定 epoch 中应该调用哪个计划程序。
lr_scheduler.ReduceLROnPlateau	当指标停止改进时降低学习率。
lr_scheduler.CyclicLR	根据循环学习率策略 (CLR) 设置每个参数组的学习率。
lr_scheduler.OneCycleLR	根据 1cycle 学习率策略设置每个参数组的学习率。
lr_scheduler.CosineAnnealingWarmRestarts	Set the learning rate of each parameter group using a cosine annealing schedule, where $\eta_{max}$ is set to the initial lr, $T_{cur}$ is the number of epochs since the last restart and $T_{i}$ is the number of epochs between two warm restarts in 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 已在 Averaging Weights Leads to Wider Optima and Better Generalization 中提出。

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

构建平均模型

AveragedModel 类用于计算 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))

衰减是一个介于 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 是一个学习率计划程序，它将学习率退火到一个固定值，然后保持不变。例如，以下代码创建一个计划程序，在每个参数组中将学习率从其初始值线性退火到 0.05，持续 5 个 epoch

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

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

处理批处理规范化

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

>>> 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 个 epoch，然后切换到 SWA 学习率时间表，并开始在 epoch 160 处收集参数的 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()
>>>
>>> # Update bn statistics for the swa_model at the end
>>> torch.optim.swa_utils.update_bn(loader, swa_model)
>>> # Use swa_model to make predictions on test data
>>> 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)
>>>
>>> # Update bn statistics for the ema_model at the end
>>> torch.optim.swa_utils.update_bn(loader, ema_model)
>>> # Use ema_model to make predictions on test data
>>> preds = ema_model(test_input)