阿姆斯格拉德

这是 PyTorch 对《亚当与超越的融合》一文的实现。

我们将其作为我们的 Adam 优化器实现的扩展。它自身的实现非常小,因为它与亚当非常相似。

我们还实现了本文中描述的合成示例,其中亚当未能收敛。

18from typing import Dict
19
20import torch
21from torch import nn
22
23from labml_nn.optimizers import WeightDecay
24from labml_nn.optimizers.adam import Adam

amsGrad 优化器

这个类是从中定义的 Adam 优化器扩展而来的adam.py 。Adam 优化器正在扩展中GenericAdaptiveOptimizer 定义的类__init__.py

27class AMSGrad(Adam):

初始化优化器

  • params 是参数列表
  • lr 是学习率
  • betas 是 (,) 的元组
  • eps基于optimized_update
  • weight_decay 是在中WeightDecay 定义的类的实例 __init__.py
  • “optimized_update” 是一个标志,在添加后是否要优化第二个时刻的偏差校正
  • amsgrad 是一个标志,指示是使用 AmsGrad 还是回退到普通的 Adam
  • defaults 是组值的默认字典。当你想扩展类时,这很有用Adam
35    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-16,
36                 weight_decay: WeightDecay = WeightDecay(),
37                 optimized_update: bool = True,
38                 amsgrad=True, defaults=None):
53        defaults = {} if defaults is None else defaults
54        defaults.update(dict(amsgrad=amsgrad))
55
56        super().__init__(params, lr, betas, eps, weight_decay, optimized_update, defaults)

初始化参数状态

  • state 是参数(张量)的优化器状态
  • group 存储参数组的优化程序属性
  • param 是参数张量
58    def init_state(self, state: Dict[str, any], group: Dict[str, any], param: nn.Parameter):

我们正在扩展init_state 的 Call of Adam 优化器

68        super().init_state(state, group, param)

如果 famsgrad lagTrue 用于此参数组,则我们保持梯度平方指数移动平均线的最大值

72        if group['amsgrad']:
73            state['max_exp_avg_sq'] = torch.zeros_like(param, memory_format=torch.preserve_format)

计算和和

  • state 是参数(张量)的优化器状态
  • group 存储参数组的优化程序属性
  • grad 是参数的当前梯度张量
75    def get_mv(self, state: Dict[str, any], group: Dict[str, any], grad: torch.Tensor):

Adam 那里得

85        m, v = super().get_mv(state, group, grad)

如果此参数组正在使用amsgrad

88        if group['amsgrad']:

得到

🗒 本文使用了这个符号,我们在这里不使用这种符号,因为它与亚当对偏差校正指数移动平均线使用相同的符号混淆了。

94            v_max = state['max_exp_avg_sq']

计算

🤔 我觉得你应该取/保持偏差校正的平方梯度的第二个指数平均值的最大值。但这也是在 PyTorch 中实现它的方式。我想这并不重要,因为偏差校正只会增加值,而且只会在训练的最初几个步骤中产生实际差异。

103            torch.maximum(v_max, v, out=v_max)
104
105            return m, v_max
106        else:

如果参数组未使用,则回退到 Adamamsgrad

108            return m, v

合成实验

这是论文中描述的合成实验,它显示了亚当失败的情景。

本文(和亚当)将优化问题描述为最小化函数相对于参数的预期值。在随机训练设置中,我们无法掌握它自身的函数;也就是说,当你优化时,NN 将是整批数据的函数。我们实际评估的是一个小批量,所以实际的功能是随机指标的实现。这就是我们谈论预期值的原因。因此,让函数实现适用于训练的每个时间步。

我们将优化器的性能作为遗憾来衡量,其中是时间步的参数是最佳的最小化的参数

现在让我们来定义综合问题,

在哪里。最佳的解决方案是

这段代码将尝试运行亚当阿姆斯格拉德来解决这个问题。

111def _synthetic_experiment(is_adam: bool):

定义参数

153    x = nn.Parameter(torch.tensor([.0]))

最佳,

155    x_star = nn.Parameter(torch.tensor([-1]), requires_grad=False)

157    def func(t: int, x_: nn.Parameter):
161        if t % 101 == 1:
162            return (1010 * x_).sum()
163        else:
164            return (-10 * x_).sum()

初始化相关的优化器

167    if is_adam:
168        optimizer = Adam([x], lr=1e-2, betas=(0.9, 0.99))
169    else:
170        optimizer = AMSGrad([x], lr=1e-2, betas=(0.9, 0.99))

172    total_regret = 0
173
174    from labml import monit, tracker, experiment

创建实验以记录结果

177    with experiment.record(name='synthetic', comment='Adam' if is_adam else 'AMSGrad'):

跑步跑

179        for step in monit.loop(10_000_000):

181            regret = func(step, x) - func(step, x_star)

183            total_regret += regret.item()

每 1000 步跟踪一次结果

185            if (step + 1) % 1000 == 0:
186                tracker.save(loss=regret, x=x, regret=total_regret / (step + 1))

计算梯度

188            regret.backward()

优化

190            optimizer.step()

渐变清晰

192            optimizer.zero_grad()

请确保

195            x.data.clamp_(-1., +1.)
196
197
198if __name__ == '__main__':

运行合成实验的是亚当。你可以看到亚当聚集在

201    _synthetic_experiment(True)

amsGrad 运行合成实验你可以看到 amsGrad 会聚到真正的最优值

204    _synthetic_experiment(False)