NSDT工具推荐: Three.js AI纹理开发包 - YOLO合成数据生成器 - GLTF/GLB在线编辑 - 3D模型格式在线转换 - 可编程3D场景编辑器 - REVIT导出3D模型插件 - 3D模型语义搜索引擎 - AI模型在线查看 - Three.js虚拟轴心开发包 - 3D模型在线减面 - STL模型在线切割 - 3D道路快速建模
本文将介绍如何在 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翻译整理,转载请标明出处