[feat] get_lr 增加可选 warmup_frac 参数(默认关闭,行为逐值不变)
这是 #862 的收窄重提。#862 同时改了 5 个训练脚本的 argparse,面太宽; 这一版只改 trainer_utils.py 里的 get_lr 本身,不动任何训练脚本、不加任何命令行参数。 get_lr(current_step, total_steps, lr) 在 current_step=0 时返回满 lr,没有 warmup。 需要 warmup 的使用者(例如从随机初始化跑较大 batch 的预训练)现在可以传第四个 参数;不传则走原来那一行表达式。 在 total_steps ∈ {100, 1000, 7339, 39634, 150000} 上取 588 个采样点, warmup_frac=0 与原实现的最大差异是 0.000e+00。pull/866/head
parent
1e6e909f88
commit
f81a948b52
|
|
@ -39,7 +39,15 @@ def Logger(content):
|
||||||
print(content)
|
print(content)
|
||||||
|
|
||||||
|
|
||||||
def get_lr(current_step, total_steps, lr):
|
def get_lr(current_step, total_steps, lr, warmup_frac=0.0):
|
||||||
|
# warmup_frac=0(默认)时与原实现逐值相同。需要 warmup 的场景可以传入一个比例,
|
||||||
|
# 前 warmup_frac*total_steps 步把 lr 线性升到满值,其后在剩余步数上走原来的余弦。
|
||||||
|
if warmup_frac > 0:
|
||||||
|
warmup_steps = max(1, int(warmup_frac * total_steps))
|
||||||
|
if current_step < warmup_steps:
|
||||||
|
return lr * (current_step + 1) / warmup_steps
|
||||||
|
current_step -= warmup_steps
|
||||||
|
total_steps = max(1, total_steps - warmup_steps)
|
||||||
return lr*(0.1 + 0.45*(1 + math.cos(math.pi * current_step / total_steps)))
|
return lr*(0.1 + 0.45*(1 + math.cos(math.pi * current_step / total_steps)))
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue