Pytorch自定义优化器

本文将介绍如何在 PyTorch 中编写自己的优化器。

optimizer = MySOTAOptimizer(my_model.parameters(), lr=0.001)
for epoch in epochs:
	for batch in epoch:
		outputs = my_model(batch)
		loss  = loss_fn(outputs, true_values)
		loss.backward()
		optimizer.step()

PyTorch 的伟大之处在于它附带了一个很棒的标准优化器库，可以满足你几乎所有的机器学习需求。然而，有时你会发现你需要一些更专业的东西。

也许你编写了自己的优化算法，该算法特别适合你正在处理的问题类型，或者你可能正在寻求从最近发表的一篇研究论文中实现一个优化器，该论文尚未进入 PyTorch 标准库。无论你的特定用例是什么，PyTorch 都允许你快速轻松地编写优化器，只要稍微了解其内部原理即可。让我们深入了解一下。

1、子类化 Pytorch 优化器类

PyTorch 中的所有优化器都需要继承自 torch.optim.Optimizer。这是处理所有通用优化机制的基类。在此类中，需要重写两个主要方法：__init__ 和 step。让我们看看它是如何完成的。

2、init方法

可以在 __init__ 方法中设置优化器的所有配置设置。 __init__方法必须采用 params 参数，该参数指定将被优化的可迭代参数。这个可迭代对象必须具有确定性的顺序 - 优化器的用户不应该传入字典或集合之类的东西。通常会给出 torch.Tensor 对象的列表。

将在 __init__ 方法中指定的其他典型参数包括 lr、weight_decays、基于 Adam 的优化器的 beta参数等。

__init__ 方法还应该对传入的参数执行一些基本检查。例如，如果提供的学习率为负，则应引发异常。

除了 params 之外，Optimizer 基类在初始化时还需要一个名为 defaults 的参数。这应该是一个将参数名称映射到其默认值的字典。它可以根据优化器类的 init 方法中收集的 kwarg 参数构建。这在接下来的内容中很重要。

__init__ 方法的最后一步是调用 Optimizer 基类。这是通过使用以下通用签名调用 super() 来执行的。

super(YourOptimizerName, self).__init__(params, defaults)

3、从零开始实现自己的优化器

让我们使用 HuggingFace pytorch-transformers NLP 库中的示例来研究和强化上述方法。他们实现了 BERT 论文中权重衰减 Adam 优化器的 PyTorch 版本。首先我们来看看类定义和 __init__ 方法。这里把两者结合起来。

可以看到 __init__ 方法完成了上面列出的所有基本要求。它对所有提供的 kwargs 的有效性进行基本检查，如果不满足则引发异常。它还根据这些所需参数构造默认值字典。最后，调用 super() 方法以使用提供的 params 和 defaults 初始化 Optimizer 基类。

4、Step()方法

真正的魔力发生在step() 方法中。这是优化器逻辑在提供的参数上实现和制定的地方。让我们看看这是如何发生的。

step(self,closure=None)中首先要注意的是closure关键字参数的存在。如果你查阅 PyTorch 文档，会发现闭包是一个可选的可调用函数，允许你在多个时间步重新评估损失。这对于大多数优化器来说是不必要的，但在共轭梯度和 LBFGS 等少数优化器中使用。根据文档，“闭包应该清除梯度，计算损失，然后返回它”。我们就这样吧，因为 AdamW 优化器不需要闭包。

关于 AdamW 步骤函数，你会注意到的下一件事是它迭代称为 param_groups 的东西。优化器的 param_groups 是一个字典列表，它提供了一种将模型参数分解为单独组件以进行优化的简单方法。它允许模型的训练者将模型参数分割成单独的单元，然后可以在不同的时间和不同的设置下进行优化。多个 param_group 的一种用途是使用不同的学习率等来训练网络的不同层。另一个突出的用例出现在迁移学习中。在微调预训练网络时，你可能需要逐渐解冻层，并随着微调的进行将它们添加到优化过程中。为此，param_groups 至关重要。

下面是 PyTorch 文档中给出的示例，其中为 SGD 指定了 param_groups，以便单独调整分类器的不同层。

现在我们已经介绍了 PyTorch 内部的一些特定内容，让我们开始讨论算法。这是最初提出 AdamW 算法的论文的链接。这里是论文中提议的更新规则的屏幕截图。

让我们通过源代码逐行浏览一下。首先，我们有循环

for p in group['params']

这里没什么神秘的。对于每个参数组，我们都会迭代该组内的参数。下一个。

if p.grad is None:
	continue
grad = p.grad.data
if grad.is_sparse:
	raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')

这也是很简单的事情。如果当前参数没有梯度，我们就跳过它。接下来，我们通过访问 p.grad.data 获取梯度的实际普通 Tensor 对象。最后，如果张量是稀疏的，我们会引发错误，因为我们不会考虑对稀疏对象实现这一点。

接下来，我们使用以下命令访问当前优化器状态

state = self.state[p]

在 PyTorch 优化器中，状态只是与优化器关联的字典，它保存所有参数的当前配置。如果这是我们第一次访问给定参数的状态，那么我们设置以下默认值

if len(state) == 0:
	state['step'] = 0
	# Exponential moving average of gradient values
	state['exp_avg'] = torch.zeros_like(p.data)
	# Exponential moving average of squared gradient values
	state['exp_avg_sq'] = torch.zeros_like(p.data)

显然，我们从步骤 0 开始，同时将指数平均值和指数平方平均值参数归零，这两个参数都是梯度张量的形状。

exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
beta1, beta2 = group['betas']

state['step'] += 1

接下来，我们从状态字典中收集将用于更新计算的参数。我们还增加当前步骤。

现在，我们开始实际的更新。这是代码。

# Decay the first and second moment running average coefficient
# In-place operations to update the averages at the same time
exp_avg.mul_(beta1).add_(1.0 - beta1, grad)
exp_avg_sq.mul_(beta2).addcmul_(1.0 - beta2, grad, grad)
denom = exp_avg_sq.sqrt().add_(group['eps'])

step_size = group['lr']
if group['correct_bias']:  # No bias correction for Bert
	bias_correction1 = 1.0 - beta1 ** state['step']
	bias_correction2 = 1.0 - beta2 ** state['step']
	step_size = step_size * math.sqrt(bias_correction2) / bias_correction1

p.data.addcdiv_(-step_size, exp_avg, denom)

# Just adding the square of the weights to the loss function is *not*
# the correct way of using L2 regularization/weight decay with Adam,
# since that will interact with the m and v parameters in strange ways.
#
# Instead we want to decay the weights in a manner that doesn't interact
# with the m/v parameters. This is equivalent to adding the square
# of the weights to the loss with plain (non-momentum) SGD.
# Add weight decay at the end (fixed version)
if group['weight_decay'] > 0.0:
	p.data.add_(-group['lr'] * group['weight_decay'], p.data)

上面的代码对应于论文中算法实现中的方程6-12。遵循数学应该很容易。我想仔细研究一下内置的张量方法，它允许我们进行就地计算。

你可能不知道的 PyTorch 的一个很好的、相对隐藏的功能是你可以访问任何标准 PyTorch 函数，例如 torch.add()、torch.mul() 等通过在方法名称后附加 _ 直接对张量进行就地操作。因此，仔细观察第一次更新，我们发现我们可以快速将其计算为

exp_avg.mul_(beta1).add_(1.0 - beta1, grad)

而不是

torch.mul(beta1, torch.add(1.0 - beta1, grad))

当然，这里使用了一些你可能不熟悉的特殊操作，例如Tensor.addcmul_和Tensor.addcdiv_。这将获取输入并将其分别添加到后两个输入的乘积或股息中。如果需要更深入地了解可在张量对象上执行的各种操作，我强烈建议你查看这篇文章。

你还会看到在计算最终结果的最后一行中访问了学习率。然后，该损失将被返还。

而且……就是这样！构建自己的优化器就这么简单。当然，你需要首先设计自己的优化算法，这可能有点棘手;)。我把那个留给你。

特别感谢 Hugging Face 的作者在 PyTorch 中实现 AdamW 优化器。

原文链接：WRITING YOUR OWN OPTIMIZERS IN PYTORCH

BimAnt翻译整理，转载请标明出处