torch.optim

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

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

如何使用优化器

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

构建它

要构建一个 Optimizer，您必须为其提供一个可迭代对象，其中包含要优化的参数（都应该是 Parameter）或命名参数（(str, Parameter) 的元组）。然后，您可以指定特定于优化器的选项，例如学习率、权重衰减等。

示例

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

命名参数示例

optimizer = optim.SGD(model.named_parameters(), lr=0.01, momentum=0.9)
optimizer = optim.Adam([('layer0', var1), ('layer1', 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)

optim.SGD([
                {'params': model.base.named_parameters(), 'lr': 1e-2},
                {'params': model.classifier.named_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 的可迭代对象。指定应优化哪些 Tensor。

• defaults (Dict[str, Any]) – (dict): 一个包含优化选项默认值的 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	实现随机梯度下降（可选动量）。

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

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

一般来说，这 3 种实现的性能排序是 fused > foreach > for-loop。因此，在适用时，我们默认使用 foreach 而不是 for-loop。适用意味着 foreach 实现可用，用户没有指定任何特定于实现的 kwargs（例如，fused、foreach、differentiable），并且所有张量都是原生的。请注意，虽然 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	beta	不支持	不支持
Adam	beta	稳定	beta
AdamW	beta	稳定	beta
SparseAdam	不支持	不支持	不支持
Adamax	不支持	不支持	不支持
ASGD	不支持	不支持	不支持
LBFGS	不支持	不支持	不支持
NAdam	不支持	不支持	不支持
RAdam	不支持	不支持	不支持
RMSprop	不支持	不支持	不支持
Rprop	不支持	不支持	不支持
SGD	beta	beta	beta

如何调整学习率

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 以 BC 破坏的方式更改了此行为。如果您在优化器更新之前（调用 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	使用余弦退火调度设置每个参数组的学习率。

如何利用命名参数加载优化器状态字典

函数 load_state_dict() 会存储加载的状态字典中的可选 param_names 内容（如果存在）。但是，加载优化器状态的过程不受影响，因为参数的顺序对于保持兼容性很重要（在顺序不同的情况下）。为了利用加载的状态字典中加载的参数名称，需要根据所需的行为实现自定义的 register_load_state_dict_pre_hook。

例如，当模型架构发生更改，但权重和优化器状态需要保持不变时，这可能很有用。以下示例演示了如何实现此自定义。

示例

class OneLayerModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(3, 4)

    def forward(self, x):
        return self.fc(x)

model = OneLayerModel()
optimizer = optim.SGD(model.named_parameters(), lr=0.01, momentum=0.9)
# training..
torch.save(optimizer.state_dict(), PATH)

假设 model 实现了专家模型 (MoE)，并且我们想要复制它并为两个专家恢复训练，这两个专家都以与 fc 层相同的方式初始化。对于以下 model2，我们创建了两个与 fc 相同的层，并通过将模型权重和优化器状态从 model 加载到 model2 的 fc1 和 fc2 中（并进行相应的调整）来恢复训练

class TwoLayerModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(3, 4)
        self.fc2 = nn.Linear(3, 4)

    def forward(self, x):
        return (self.fc1(x) + self.fc2(x)) / 2

model2 = TwoLayerModel()
# adapt and load model weights..
optimizer2 = optim.SGD(model2.named_parameters(), lr=0.01, momentum=0.9)

为了使用先前优化器的状态字典加载 optimizer2 的状态字典，以便 fc1 和 fc2 都将使用 fc 优化器状态的副本初始化（以从 fc 恢复每个层的训练），我们可以使用以下钩子

def adapt_state_dict_ids(optimizer, state_dict):
    adapted_state_dict = deepcopy(optimizer.state_dict())
    # Copy setup parameters (lr, weight_decay, etc.), in case they differ in the loaded state dict.
    for k, v in state_dict['param_groups'][0].items():
        if k not in ['params', 'param_names']:
            adapted_state_dict['param_groups'][0][k] = v

    lookup_dict = {
        'fc1.weight': 'fc.weight',
        'fc1.bias': 'fc.bias',
        'fc2.weight': 'fc.weight',
        'fc2.bias': 'fc.bias'
    }
    clone_deepcopy = lambda d: {k: (v.clone() if isinstance(v, torch.Tensor) else deepcopy(v)) for k, v in d.items()}
    for param_id, param_name in zip(
            optimizer.state_dict()['param_groups'][0]['params'],
            optimizer.state_dict()['param_groups'][0]['param_names']):
        name_in_loaded = lookup_dict[param_name]
        index_in_loaded_list = state_dict['param_groups'][0]['param_names'].index(name_in_loaded)
        id_in_loaded = state_dict['param_groups'][0]['params'][index_in_loaded_list]
        # Copy the state of the corresponding parameter
        if id_in_loaded in state_dict['state']:
            adapted_state_dict['state'][param_id] = clone_deepcopy(state_dict['state'][id_in_loaded])

    return adapted_state_dict

optimizer2.register_load_state_dict_pre_hook(adapt_state_dict_ids)
optimizer2.load_state_dict(torch.load(PATH)) # The previous optimizer saved state_dict

这确保了在模型加载期间将使用针对 model2 层进行调整的 state_dict 以及正确的状态。请注意，此代码专门为此示例设计（例如，假设单个参数组），其他情况可能需要不同的调整。

以下示例展示了当模型结构更改时，如何在加载的 state dict 中处理缺失的参数。Model_bypass 添加了一个新的 bypass 层，该层在原始 Model1 中不存在。为了恢复训练，使用了自定义的 adapt_state_dict_missing_param 钩子来调整优化器的 state_dict，确保现有参数被正确映射，而缺失的参数（如 bypass 层）保持不变（如本示例中初始化）。这种方法能够在模型更改的情况下平滑加载和恢复优化器状态。新的 bypass 层将从头开始训练

class Model1(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(5, 5)

    def forward(self, x):
        return self.fc(x) + x


model = Model1()
optimizer = optim.SGD(model.named_parameters(), lr=0.01, momentum=0.9)
# training..
torch.save(optimizer.state_dict(), PATH)

class Model_bypass(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(5, 5)
        self.bypass = nn.Linear(5, 5, bias=False)
        torch.nn.init.eye_(self.bypass.weight)

    def forward(self, x):
        return self.fc(x) + self.bypass(x)

model2 = Model_bypass()
optimizer2 = optim.SGD(model2.named_parameters(), lr=0.01, momentum=0.9)

def adapt_state_dict_missing_param(optimizer, state_dict):
    adapted_state_dict = deepcopy(optimizer.state_dict())
    # Copy setup parameters (lr, weight_decay, etc.), in case they differ in the loaded state dict.
    for k, v in state_dict['param_groups'][0].items():
        if k not in ['params', 'param_names']:
            adapted_state_dict['param_groups'][0][k] = v

    lookup_dict = {
        'fc.weight': 'fc.weight',
        'fc.bias': 'fc.bias',
        'bypass.weight': None,
    }

    clone_deepcopy = lambda d: {k: (v.clone() if isinstance(v, torch.Tensor) else deepcopy(v)) for k, v in d.items()}
    for param_id, param_name in zip(
            optimizer.state_dict()['param_groups'][0]['params'],
            optimizer.state_dict()['param_groups'][0]['param_names']):
        name_in_loaded = lookup_dict[param_name]
        if name_in_loaded in state_dict['param_groups'][0]['param_names']:
            index_in_loaded_list = state_dict['param_groups'][0]['param_names'].index(name_in_loaded)
            id_in_loaded = state_dict['param_groups'][0]['params'][index_in_loaded_list]
            # Copy the state of the corresponding parameter
            if id_in_loaded in state_dict['state']:
                adapted_state_dict['state'][param_id] = clone_deepcopy(state_dict['state'][id_in_loaded])

    return adapted_state_dict

optimizer2.register_load_state_dict_pre_hook(adapt_state_dict_ids)
optimizer2.load_state_dict(torch.load(PATH)) # The previous optimizer saved state_dict

作为第三个示例，与其按照参数的顺序加载状态（默认方法），不如使用此钩子根据参数名称加载

def names_matching(optimizer, state_dict):
    assert len(state_dict['param_groups']) == len(optimizer.state_dict()['param_groups'])
    adapted_state_dict = deepcopy(optimizer.state_dict())
    for g_ind in range(len(state_dict['param_groups'])):
        assert len(state_dict['param_groups'][g_ind]['params']) == len(
            optimizer.state_dict()['param_groups'][g_ind]['params'])

        for k, v in state_dict['param_groups'][g_ind].items():
            if k not in ['params', 'param_names']:
                adapted_state_dict['param_groups'][g_ind][k] = v

        for param_id, param_name in zip(
                optimizer.state_dict()['param_groups'][g_ind]['params'],
                optimizer.state_dict()['param_groups'][g_ind]['param_names']):
            index_in_loaded_list = state_dict['param_groups'][g_ind]['param_names'].index(param_name)
            id_in_loaded = state_dict['param_groups'][g_ind]['params'][index_in_loaded_list]
            # Copy the state of the corresponding parameter
            if id_in_loaded in state_dict['state']:
                adapted_state_dict['state'][param_id] = deepcopy(state_dict['state'][id_in_loaded])

    return adapted_state_dict

权重平均 (SWA 和 EMA)

torch.optim.swa_utils.AveragedModel 实现了随机权重平均 (SWA) 和指数移动平均 (EMA)，torch.optim.swa_utils.SWALR 实现了 SWA 学习率调度器，torch.optim.swa_utils.update_bn() 是一个实用函数，用于在训练结束时更新 SWA/EMA 批归一化统计信息。

SWA 已在 Averaging Weights Leads to Wider Optima and Better Generalization 中提出。

EMA 是一种广为人知的技术，通过减少所需的权重更新次数来减少训练时间。它是 Polyak averaging 的变体，但使用指数权重而不是迭代中相等的权重。

构建平均模型

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

Decay 是一个介于 0 和 1 之间的参数，用于控制平均参数衰减的速度。如果未提供给 torch.optim.swa_utils.get_ema_multi_avg_fn()，则默认值为 0.999。Decay 值应接近 1.0，因为较小的值可能会导致优化收敛问题。

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 将跟踪模型参数的运行平均值。要更新这些平均值，您应该在 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 学习率 schedules

通常，在 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 学习率 schedule 并开始收集参数的 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()。