学习率调整
import math
def adjust_learning_rate(optimizer, epoch, args):
"""Decay the learning rate with half-cycle cosine after warmup"""
if epoch < args.warmup_epochs:
lr = args.lr * epoch / args.warmup_epochs
else:
lr = args.min_lr + (args.lr - args.min_lr) * 0.5 * \
(1. + math.cos(math.pi * (epoch - args.warmup_epochs) / (args.epochs - args.warmup_epochs)))
for param_group in optimizer.param_groups:
if "lr_scale" in param_group:
param_group["lr"] = lr * param_group["lr_scale"]
else:
param_group["lr"] = lr
return lr
预热阶段的学习率调整:
if epoch < args.warmup_epochs:
lr = args.lr * epoch / args.warmup_epochs
如果当前 epoch
小于预热周期数 args.warmup_epochs
,则学习率 lr
随着 epoch
线性增加,从 0 增加到初始学习率 args.lr
。公式为 lr = lr * epoch / warmup_epochs
:
lr
=
lr
×
epoch
w
a
r
m
u
p
_
e
p
o
c
h
s
\text{lr} = \frac{\text{lr} \times \text{epoch}}{warmup\_epochs}
lr=warmup_epochslr×epoch
余弦退火阶段的学习率调整
else:
lr = args.min_lr + (args.lr - args.min_lr) * 0.5 * \
(1. + math.cos(math.pi * (epoch - args.warmup_epochs) / (args.epochs - args.warmup_epochs)))
如果当前 epoch
大于等于预热周期数 args.warmup_epochs
,则使用余弦退火函数来调整学习率。公式为:
lr = args.min_lr + ( args.lr − args.min_lr ) × 0.5 × ( 1 + cos ( π × epoch − args.warmup_epochs args.epochs − args.warmup_epochs ) ) \text{lr} = \text{args.min\_lr} + (\text{args.lr} - \text{args.min\_lr}) \times 0.5 \times \left(1 + \cos\left(\pi \times \frac{\text{epoch} - \text{args.warmup\_epochs}}{\text{args.epochs} - \text{args.warmup\_epochs}}\right)\right) lr=args.min_lr+(args.lr−args.min_lr)×0.5×(1+cos(π×args.epochs−args.warmup_epochsepoch−args.warmup_epochs))
更新优化器的学习率
for param_group in optimizer.param_groups:
if "lr_scale" in param_group:
param_group["lr"] = lr * param_group["lr_scale"]
else:
param_group["lr"] = lr
这个函数根据训练周期动态调整学习率,先进行线性预热,然后使用余弦退火策略。这样可以在初期稳定模型的训练,避免较大学习率带来的不稳定性,并在后期逐渐减小学习率,提高模型的收敛效果。