torch.optim

构建

要构建一个Optimizer，你需要提供一个包含参数（所有参数都应该是Variables）的可迭代对象。然后，你可以指定特定优化器的选项，如学习率、权重衰减等。

示例：

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

参数选项

Optimizer 还支持为每个参数指定不同的选项。为此，不要传递 Variable 的可迭代对象，而是传入一个字典的可迭代对象。每个字典定义一个单独的参数组，并应包含一个 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() 方法，用于更新参数。此方法可以有两种用法：

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

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)

警告

在 PyTorch 1.1.0 之前，学习率调度器需要在优化器更新之前被调用；而从 1.1.0 版本开始，这一行为以破坏向后兼容性的方式进行了更改。如果你在调用 optimizer.step() 之前调用了学习率调度器（即调用 scheduler.step()），这将导致跳过学习率调度的第一个值。因此，如果你在升级到 PyTorch 1.1.0 后无法重现结果，请检查是否在错误的时间调用了 scheduler.step()。

构建平均模型

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

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

>>> averaged_model = AveragedModel(model)

可以通过以下方式指定 multi_avg_fn 参数来构建 EMA 模型：

>>> 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）公式应用到权重上。

其中 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个epochs内完成：

>>> 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 应用到数据加载器中的每一个元素，并计算模型中每个批处理规范化层的激活统计信息。

综合介绍：SWA

在下面的例子中，swa_model 是一个累积权重平均值的SWA模型。我们总共训练300个 epochs，并从第160个 epoch 开始切换到 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)

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) – 用于更新批量归一化统计信息的模型。

• device (torch.device, 可选) – 如果设置了设备，数据将在传递给 model 之前被转移到指定的 device。

示例

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

注意

update_bn 工具假设 loader 中的每个数据批次要么是一个张量，要么是张量的列表或元组。如果是后者，则应调用 model.forward() 来处理对应于数据批次的第一个元素。