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)

所有优化器的基类。

警告

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

参数

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

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

Optimizer.add_param_group	将参数组添加到 Optimizer 的 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 预钩子，它将在调用 state_dict() 之前被调用。
Optimizer.register_state_dict_post_hook	注册一个 state dict 后钩子，它将在调用 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	实现随机梯度下降（可选地使用动量）。

我们的许多算法都有针对性能、可读性和/或通用性优化的各种实现，因此我们尝试在用户未指定任何特定实现的关键字参数（例如，fused、foreach、可微分）的情况下，默认使用当前设备上通常最快的实现。

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

一般来说，3 种实现的性能排序是 fused > foreach > for-loop。因此，在适用时，我们默认使用 foreach 而不是 for-loop。适用是指 foreach 实现可用，用户未指定任何特定于实现的关键字参数（例如，fused、foreach、可微分），并且所有张量都是本地的。请注意，虽然 fused 应该比 foreach 更快，但这些实现更新，我们希望在将所有地方都切换之前，先让它们更多地烘焙。我们在下面的第二个表中总结了每种实现的稳定性状态，欢迎您尝试一下！

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

算法	默认	有 foreach 吗？	有 fused 吗？
Adadelta	foreach	是	否
Adafactor	for-loop	否	否
Adagrad	foreach	是	是（仅限 cpu）
Adam	foreach	是	是
AdamW	foreach	是	是
SparseAdam	for-loop	否	否
Adamax	foreach	是	否
ASGD	foreach	是	否
LBFGS	for-loop	否	否
NAdam	foreach	是	否
RAdam	foreach	是	否
RMSprop	foreach	是	否
Rprop	foreach	是	否
SGD	foreach	是	是

以下表格显示了 fused 实现的稳定性状态

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

如何调整学习率

torch.optim.lr_scheduler.LRScheduler 提供了几种方法来根据 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.LRScheduler	在优化过程中调整学习率。
lr_scheduler.LambdaLR	设置初始学习率。
lr_scheduler.MultiplicativeLR	将每个参数组的学习率乘以指定函数中给出的因子。
lr_scheduler.StepLR	每隔 step_size 个 epoch 将每个参数组的学习率衰减 gamma。
lr_scheduler.MultiStepLR	当 epoch 数达到某个里程碑时，将每个参数组的学习率衰减 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 是一个学习率调度器，它将学习率退火到一个固定值，然后保持不变。 例如，以下代码创建了一个调度器，它在每个参数组中将学习率从其初始值线性退火到 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，并在 epoch 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 个 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)

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 层的激活统计信息。

参数

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

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

• device (torch.device, optional) – 如果设置，数据将在传递到 model 之前被传输到 device。

示例

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

注意

该 update_bn 实用程序假设 loader 中的每个数据批次都是张量或张量列表或元组；在后一种情况下，假设应该在对应于数据批次的列表或元组的第一个元素上调用 model.forward()。