#

阿姆斯格拉德

我们将其作为我们的 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 是 ( $β_{1}$ , $β_{2}$ ) 的元组
eps 是 $\overset{ϵ}{^}$ 或 $ϵ$ 基于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 是参数张量 $θ_{t - 1}$

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)

#

计算 $m_{t}$ 和和 $v_{t}$ 或 $m a x (v_{1}, v_{2}, ..., v_{t - 1}, v_{t})$

state 是参数（张量）的优化器状态
group 存储参数组的优化程序属性
grad 是参数的当前梯 $g_{t}$ 度张量 $θ_{t - 1}$

75    def get_mv(self, state: Dict[str, any], group: Dict[str, any], grad: torch.Tensor):

#

$v_{t}$ 从 Adam 那里得 $m_{t}$ 到

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

#

如果此参数组正在使用amsgrad

88        if group['amsgrad']:

#

得到 $max (v_{1}, v_{2}, ..., v_{t - 1})$ 。

🗒 本文使用了这个符号 $\overset{v}{^}_{t}$ ，我们在这里不使用这种符号，因为它与亚当对偏差校正指数移动平均线使用相同的符号混淆了。

94            v_max = state['max_exp_avg_sq']

#

计算 $m a x (v_{1}, v_{2}, ..., v_{t - 1}, v_{t})$ 。

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

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

#

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

108            return m, v

#

合成实验

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

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

我们将优化器的性能作为遗憾来衡量， $R (T) = t = 1 \sum T [f_{t} (θ_{t}) - f_{t} (θ^{*})]$ 其中 $θ_{t}$ 是时间步的参数 $t$ ， $θ^{*}$ 是最佳的最小化的参数 $E [f (θ)]$ 。

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

f_{t} (x) = {1010 x, - 10 x, f or t mod 101 = 1 o t herwise

在哪里 $- 1 \leq x \leq + 1$ 。最佳的解决方案是 $x = - 1$ 。

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

111def _synthetic_experiment(is_adam: bool):

#

定义 $x$ 参数

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

#

最佳， $x^{*} = - 1$

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

#

$f_{t} (x)$

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

#

$R (T)$

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

#

跑步跑 $1 0^{7}$ 步

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

#

$f_{t} (θ_{t}) - f_{t} (θ^{*})$

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

#

$R (T) = \sum_{t = 1}^{T} [f_{t} (θ_{t}) - f_{t} (θ^{*})]$

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

#

请确保 $- 1 \leq x \leq + 1$

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

#

运行合成实验的是亚当。你可以看到亚当聚集在 $x = + 1$

201    _synthetic_experiment(True)

#

在 amsGrad 运行合成实验你可以看到 amsGrad 会聚到真正的最优值 $x = - 1$

204    _synthetic_experiment(False)