home optimizers

Github Twitter

View code on Github

#

带有预热和余弦衰减的 Adam Optimizer

这扩展了 AmsGrad 优化器并增加了预热阶段。

11import math
12from typing import Dict
13
14from labml_nn.optimizers import WeightDecay
15from labml_nn.optimizers.amsgrad import AMSGrad
#

带有预热和余弦衰减的 Adam Optimizer

这个类是从中定义的 AmsGrad 优化器扩展而来的amsgrad.py 。

18class AdamWarmupCosineDecay(AMSGrad):
#

初始化优化器

  • params 是参数列表
  • lr 是学习率
  • betas 是 (,) 的元组
  • eps 是或基于optimized_update
  • weight_decay 是在中WeightDecay 定义的类的实例 __init__.py
  • “optimized_update” 是一个标志，在添加后是否要优化第二个时刻的偏差校正
  • amsgrad 是一个标志，指示是使用 AmsGrad 还是回退到普通的 Adam
  • warmup 预热步数
  • total_steps 总步数。此时余弦衰减达到 0，但保持在 10%，lr 因为我们得
    • 到
  • defaults 是组值的默认字典。当你想扩展类时，这很有用AdamWarmup 。
27    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-16,
28                 weight_decay: WeightDecay = WeightDecay(),
29                 optimized_update: bool = True,
30                 amsgrad=False, warmup=0, total_steps=1e10, defaults=None):
# 
49        defaults = {} if defaults is None else defaults
50        defaults.update(dict(warmup=warmup, total_steps=total_steps))
51        super().__init__(params, lr, betas, eps, weight_decay, optimized_update, amsgrad, defaults)
#

获取学习率

其中是预热步骤的数量。

53    def get_lr(self, state: Dict[str, any], group: Dict[str, any]):
#

如果我们处于热身阶段

61        if group['warmup'] > state['step']:
#

学习率从线性增加到

63            return 1e-8 + state['step'] * group['lr'] / group['warmup']
64        else:
#

持续的学习速率

66            progress = (state['step'] - group['warmup']) / max(1, group['total_steps'] - group['warmup'])
67            return group['lr'] * max(0.1, 0.5 * (1.0 + math.cos(math.pi * progress)))
#

绘制不同预热和模型大小的学习速率

Plot of learning rate

70def _test_lr():
# 
76    import matplotlib.pyplot as plt
77    import numpy as np
78    from torch import nn
79
80    model = nn.Linear(10, 10)
81    opt = AdamWarmupCosineDecay(model.parameters(), warmup=5000, lr=1e-4, total_steps=4e6)
82    steps = 20_000
83    plt.plot(np.arange(1, steps), [opt.get_lr({'step': i}, opt.defaults) for i in range(1, steps)])
84    plt.legend(["5000:4e6", "5000:2e6", "5000:1e6"])
85    plt.title("Learning Rate")
86    plt.show()
87
88    steps = int(6e6)
89    step_size = 1000
90    plt.plot(np.arange(1, steps, step_size), [opt.get_lr({'step': i}, opt.defaults) for i in range(1, steps, step_size)])
91    plt.legend(["5000:4e6", "5000:2e6", "5000:1e6"])
92    plt.title("Learning Rate")
93    plt.show()
94
95
96if __name__ == '__main__':
97    _test_lr()
labml.ai