快捷方式

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)[source]

所有优化器的基类。

警告

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

参数
  • params (iterable) – torch.Tensordict 的可迭代对象。指定哪些张量应该被优化。

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

Optimizer.add_param_group

将参数组添加到 Optimizerparam_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(多张量)和融合。最直接的实现是使用参数的 for-loop,其中包含大量计算。for-loop 通常比我们的 foreach 实现慢,foreach 实现将参数组合成多张量并一次运行大量计算,从而节省了许多顺序内核调用。我们的一些优化器甚至具有更快的融合实现,这些实现将大量计算融合成一个内核。我们可以将 foreach 实现视为水平融合,将融合实现视为在此基础上的垂直融合。

一般来说,这 3 种实现的性能排序为:融合 > foreach > for-loop。因此,在适用时,我们默认使用 foreach 而不是 for-loop。适用意味着 foreach 实现可用,用户未指定任何特定于实现的 kwargs(例如,fused、foreach、differentiable),并且所有张量都是本地的并且位于 CUDA 上。请注意,虽然 fused 应该比 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 提供了几种方法来根据 epochs 数调整学习率。 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 以破坏 BC 的方式更改了此行为。如果你在优化器更新(调用 optimizer.step())之前调用学习率调度器(调用 scheduler.step()),这将跳过学习率调度表的第一个值。如果你在升级到 PyTorch 1.1.0 之后无法重现结果,请检查你是否在错误的时间调用了 scheduler.step()

lr_scheduler.LambdaLR

将每个参数组的学习率设置为初始 lr 乘以给定函数。

lr_scheduler.MultiplicativeLR

将每个参数组的学习率乘以指定函数中给定的因子。

lr_scheduler.StepLR

每隔 step_size 个 epochs,将每个参数组的学习率降低 gamma 倍。

lr_scheduler.MultiStepLR

当 epochs 数达到某个里程碑时,将每个参数组的学习率降低 gamma 倍。

lr_scheduler.ConstantLR

将每个参数组的学习率乘以一个小常数因子,直到 epochs 数达到预定义的里程碑:total_iters。

lr_scheduler.LinearLR

通过线性地改变小乘数因子,将每个参数组的学习率降低,直到 epochs 数达到预定义的里程碑:total_iters。

lr_scheduler.ExponentialLR

每个 epoch 将每个参数组的学习率降低 gamma 倍。

lr_scheduler.PolynomialLR

使用给定 total_iters 中的多项式函数,将每个参数组的学习率降低。

lr_scheduler.CosineAnnealingLR

使用余弦退火调度设置每个参数组的学习率,其中 ηmax\eta_{max} 设置为初始 lr,而 TcurT_{cur} 是自 SGDR 上次重启以来的 epochs 数

lr_scheduler.ChainedScheduler

链接学习率调度器列表。

lr_scheduler.SequentialLR

接收预计在优化过程中按顺序调用的调度器列表,以及提供精确间隔的里程碑点,以反映在给定 epoch 时应该调用哪个调度器。

lr_scheduler.ReduceLROnPlateau

当指标停止改进时降低学习率。

lr_scheduler.CyclicLR

根据循环学习率策略 (CLR) 设置每个参数组的学习率。

lr_scheduler.OneCycleLR

根据 1cycle 学习率策略设置每个参数组的学习率。

lr_scheduler.CosineAnnealingWarmRestarts

使用余弦退火调度设置每个参数组的学习率,其中 ηmax\eta_{max} 设置为初始 lr,TcurT_{cur} 是自上次重启以来的 epochs 数,而 TiT_{i} 是 SGDR 中两次热重启之间的 epochs 数

权重平均(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平均的变体,但使用指数权重而不是跨迭代的相等权重。

构建平均模型

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方程应用于权重

Wt+1EMA=αWtEMA+(1α)WtmodelW^\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_fnmulti_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()是一个实用函数,允许在训练结束时根据给定的数据加载器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个纪元,并在第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)

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得您的问题解答

查看资源