Merge f81a948b52 into f659b55761

trainer/trainer_utils.py

@@ -39,7 +39,15 @@ def Logger(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)))