torch.optim

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

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

如何使用优化器

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

构建它

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

Example:

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() 计算梯度后，可以调用该函数。

Example:

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

optimizer.step(closure)

某些优化算法，如共轭梯度法（Conjugate Gradient）和 LBFGS，需要多次重新计算函数，因此你必须传入一个闭包（closure），以便它们能够重新计算你的模型。该闭包应该清除梯度，计算损失，并返回该损失。

Example:

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

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

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

Optimizer.add_param_group	向 Optimizer s 的 param_groups 添加一个参数组。
Optimizer.load_state_dict	加载优化器状态。
Optimizer.register_load_state_dict_pre_hook	注册一个 load_state_dict 预挂钩，它将在调用 load_state_dict() 之前被调用。它的签名应如下：.
Optimizer.register_load_state_dict_post_hook	注册一个 load_state_dict 后钩子，它将在调用 load_state_dict() 之后被调用。它的签名应如下：.
Optimizer.state_dict	返回优化器的状态作为一个 dict。
Optimizer.register_state_dict_pre_hook	注册一个状态字典预钩子，它将在调用 state_dict() 之前被调用。
Optimizer.register_state_dict_post_hook	注册一个状态字典后处理钩子，它将在调用 state_dict() 之后被调用。
Optimizer.step	执行一次优化步骤以更新参数。
Optimizer.register_step_pre_hook	注册一个优化器步骤前置钩子，该钩子将在优化器步骤之前被调用。
Optimizer.register_step_post_hook	注册一个优化器步骤后钩子，该钩子将在优化器步骤之后被调用。
Optimizer.zero_grad	重置所有优化的 torch.Tensor 的梯度。

算法

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

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

我们有三类主要的实现方式：for 循环、foreach（多张量）和融合（fused）。最直接的实现是通过参数进行 for 循环，每次处理大量计算。通常情况下，for 循环的速度比我们的 foreach 实现要慢，后者将参数组合成一个多张量，并一次性运行所有大块计算，从而节省了许多顺序内核调用。我们的一些优化器甚至具有更快的融合实现，这些实现将大块计算融合到一个内核中。我们可以将 foreach 实现视为水平方向上的融合，而 fused 实现则是在此基础上垂直方向上的进一步融合。

通常，这 3 种实现方式的性能排序为 fused > foreach > for-loop。因此，在适用的情况下，我们默认使用 foreach 而不是 for-loop。适用意味着 foreach 实现可用，用户没有指定任何特定于实现的参数（例如 fused、foreach、differentiable），并且所有张量都是原生的。请注意，虽然 fused 应该比 foreach 更快，但这些实现较新，我们希望在全面切换之前给予它们更多的测试时间。我们在下表中总结了每种实现的稳定性状态，不过欢迎您尝试！

下面是显示每种算法可用实现和默认实现的表格：

算法	默认	是否具有 foreach？	已融合？
Adadelta	foreach	是的	no
Adafactor	for-loop	no	no
Adagrad	foreach	是的	yes (仅 cpu)
Adam	foreach	是的	是的
AdamW	foreach	是的	是的
SparseAdam	for-loop	no	no
Adamax	foreach	是的	no
ASGD	foreach	是的	no
LBFGS	for-loop	no	no
NAdam	foreach	是的	no
RAdam	foreach	是的	no
RMSprop	foreach	是的	no
Rprop	foreach	是的	no
SGD	foreach	是的	是的

下表显示了融合实现的稳定性状态：

算法	CPU	CUDA	MPS
Adadelta	不受支持的	不受支持的	不受支持的
Adafactor	不受支持的	不受支持的	不受支持的
Adagrad	测试版	不受支持的	不受支持的
Adam	测试版	稳定版	测试版
AdamW	测试版	稳定版	测试版
SparseAdam	不受支持的	不受支持的	不受支持的
Adamax	不受支持的	不受支持的	不受支持的
ASGD	不受支持的	不受支持的	不受支持的
LBFGS	不受支持的	不受支持的	不受支持的
NAdam	不受支持的	不受支持的	不受支持的
RAdam	不受支持的	不受支持的	不受支持的
RMSprop	不受支持的	不受支持的	不受支持的
Rprop	不受支持的	不受支持的	不受支持的
SGD	测试版	测试版	测试版

如何调整学习率

torch.optim.lr_scheduler.LRScheduler 提供了多种方法，可根据训练轮次调整学习率。 torch.optim.lr_scheduler.ReduceLROnPlateau 允许根据某些验证指标动态降低学习率。

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

Example:

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()

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

Example:

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.LRScheduler	在优化过程中调整学习率。
lr_scheduler.LambdaLR	设置初始学习率。
lr_scheduler.MultiplicativeLR	将每个参数组的学习率乘以指定函数中给出的因子。
lr_scheduler.StepLR	每个参数组的学习率每隔 step_size 个 epoch 按 gamma 的比例衰减。
lr_scheduler.MultiStepLR	当训练轮数达到指定的里程碑之一时，将每个参数组的学习率按 gamma 进行衰减。
lr_scheduler.ConstantLR	将每个参数组的学习率乘以一个小的常数因子。
lr_scheduler.LinearLR	通过线性变化的小乘数因子衰减每个参数组的学习率。
lr_scheduler.ExponentialLR	每个参数组的学习率在每个训练周期（epoch）后按 gamma 的比例衰减。
lr_scheduler.PolynomialLR	使用给定的 total_iters 中的多项式函数衰减每个参数组的学习率。
lr_scheduler.CosineAnnealingLR	使用余弦退火计划设置每个参数组的学习率。
lr_scheduler.ChainedScheduler	将一系列学习率调度器串联起来。
lr_scheduler.SequentialLR	包含在优化过程中预期按顺序调用的调度器列表。
lr_scheduler.ReduceLROnPlateau	当某个指标停止提升时，降低学习率。
lr_scheduler.CyclicLR	根据循环学习率策略（CLR）设置每个参数组的学习率。
lr_scheduler.OneCycleLR	根据 1cycle 学习率策略设置每个参数组的学习率。
lr_scheduler.CosineAnnealingWarmRestarts	使用余弦退火计划设置每个参数组的学习率。

权重平均（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 平均 的一种变体，但在迭代中使用指数权重而不是等权重。

构建平均模型

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之间的参数，用于控制平均参数的衰减速度。如果未提供给 torch.optim.swa_utils.get_ema_multi_avg_fn()，默认值为 0.999。

torch.optim.swa_utils.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()
>>>
>>> # 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 个周期，并立即开始收集 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)

swa_utils.AveragedModel	为随机权重平均（SWA）和指数移动平均（EMA）实现平均模型。
swa_utils.SWALR	将每个参数组的学习率退火到一个固定值。

torch.optim.swa_utils.get_ema_multi_avg_fn(decay=0.999)

获取在多个参数上应用指数移动平均（EMA）的函数。

torch.optim.swa_utils.update_bn(loader, model, device=None)

更新模型中 BatchNorm 的 running_mean、running_var 缓冲区。

它对数据进行一次处理，在 loader 处估计模型中 BatchNorm 层的激活统计信息。

Parameters

• loader (torch.utils.data.DataLoader) – 用于计算激活统计的 数据集加载器。每个数据批次应该是一个张量，或者是一个列表/元组，其第一个元素是包含数据的张量。

• 模型 (torch.nn.Module) – 我们要更新其 BatchNorm 统计信息的模型。

• 设备 (torch.device, 可选) – 如果设置，数据将在传递给 device 之前传输到该设备。

示例

>>> loader, model = ...
>>> torch.optim.swa_utils.update_bn(loader, model)

注意

The update_bn 工具假定 loader 中的每个数据批次 要么是一个张量，要么是一个张量列表或元组；在后一种情况下，假定应将 model.forward() 应用于与数据批次对应的列表或元组的第一个元素。