Polynomial Decay with Linear Warm-up

BERT-style fine-tuning warms the learning rate up linearly, then decays it polynomially to an end value.

Implement poly_lr(step, total, warmup, lr_max, lr_end=0.0, power=1.0):

• step < warmup: lr_max * step / warmup
• otherwise: lr_end + (lr_max - lr_end) * (1 - (step - warmup) / (total - warmup)) ** power
• steps past total stay at lr_end