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.Tensor
或dict
的可迭代对象。指定哪些张量应该被优化。defaults (Dict[str, Any]) – (字典): 一个字典,包含优化选项的默认值(当参数组未指定时使用)。
将参数组添加到 |
|
加载优化器状态。 |
|
将优化器状态作为 |
|
执行单个优化步骤(参数更新)。 |
|
重置所有优化 |
算法¶
实现 Adadelta 算法。 |
|
实现 Adagrad 算法。 |
|
实现 Adam 算法。 |
|
实现 AdamW 算法。 |
|
SparseAdam 实现 Adam 算法的掩码版本,适用于稀疏梯度。 |
|
实现 Adamax 算法(基于无穷范数的 Adam 变体)。 |
|
实现平均随机梯度下降。 |
|
实现 L-BFGS 算法。 |
|
实现 NAdam 算法。 |
|
实现 RAdam 算法。 |
|
实现 RMSprop 算法。 |
|
实现弹性反向传播算法。 |
|
实现随机梯度下降(可选地使用动量)。 |
我们的许多算法都具有针对性能、可读性和/或通用性优化的各种实现,因此我们尝试在用户未指定任何特定实现的情况下,默认使用当前设备上最快的实现。
我们有 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? |
---|---|---|---|
foreach |
是 |
否 |
|
foreach |
是 |
否 |
|
foreach |
是 |
是 |
|
foreach |
是 |
是 |
|
for-loop |
否 |
否 |
|
foreach |
是 |
否 |
|
foreach |
是 |
否 |
|
for-loop |
否 |
否 |
|
foreach |
是 |
否 |
|
foreach |
是 |
否 |
|
foreach |
是 |
否 |
|
foreach |
是 |
否 |
|
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 乘以给定函数。 |
|
将每个参数组的学习率乘以指定函数中给定的因子。 |
|
每隔 step_size 个 epochs,将每个参数组的学习率降低 gamma 倍。 |
|
当 epochs 数达到某个里程碑时,将每个参数组的学习率降低 gamma 倍。 |
|
将每个参数组的学习率乘以一个小常数因子,直到 epochs 数达到预定义的里程碑:total_iters。 |
|
通过线性地改变小乘数因子,将每个参数组的学习率降低,直到 epochs 数达到预定义的里程碑:total_iters。 |
|
每个 epoch 将每个参数组的学习率降低 gamma 倍。 |
|
使用给定 total_iters 中的多项式函数,将每个参数组的学习率降低。 |
|
使用余弦退火调度设置每个参数组的学习率,其中 设置为初始 lr,而 是自 SGDR 上次重启以来的 epochs 数 |
|
链接学习率调度器列表。 |
|
接收预计在优化过程中按顺序调用的调度器列表,以及提供精确间隔的里程碑点,以反映在给定 epoch 时应该调用哪个调度器。 |
|
当指标停止改进时降低学习率。 |
|
根据循环学习率策略 (CLR) 设置每个参数组的学习率。 |
|
根据 1cycle 学习率策略设置每个参数组的学习率。 |
|
使用余弦退火调度设置每个参数组的学习率,其中 设置为初始 lr, 是自上次重启以来的 epochs 数,而 是 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方程应用于权重
其中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()
是一个实用函数,允许在训练结束时根据给定的数据加载器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)