"""演示21: 学习率调度 — Step / Cosine / Warm-up / Plateau。
画出 4 种经典学习率调度策略随训练步数的变化曲线。
"""
import os
import numpy as np
import matplotlib.pyplot as plt
from figure_style import figures_dir, configure
configure()
OUT = figures_dir()
os.makedirs(OUT, exist_ok=True)
steps = np.arange(1000)
base = 0.1
def step_decay(s, lr0=base, drop=0.1, every=300):
return lr0 * (drop ** (s // every))
def cosine(s, lr0=base, T=1000, lr_min=1e-4):
return lr_min + 0.5 * (lr0 - lr_min) * (1 + np.cos(np.pi * s / T))
def warmup(s, lr0=base, warm=100):
return lr0 * min(1.0, s / warm)
def warmup_cosine(s, lr0=base, warm=100, T=1000):
if s < warm:
return lr0 * s / warm
return cosine(s - warm, lr0, T - warm)
fig, ax = plt.subplots(figsize=(9, 5))
ax.plot(steps, [step_decay(s) for s in steps], label="Step Decay (每 300 步 ×0.1)")
ax.plot(steps, [cosine(s) for s in steps], label="Cosine Annealing")
ax.plot(steps, [warmup(s) for s in steps], label="Warm-up (前 100 步线性升)")
ax.plot(steps, [warmup_cosine(s) for s in steps], label="Warm-up + Cosine")
ax.set_xlabel("训练步数"); ax.set_ylabel("学习率")
ax.set_title("学习率调度: 大学习率快速前进, 小学习率精细收敛")
ax.legend(); ax.grid(alpha=0.3); ax.set_yscale("log")
fig.tight_layout()
out = os.path.join(OUT, "fig_train_lr.png")
fig.savefig(out, dpi=150)
print(f"-> {out}")
|