余弦退火学习率 CosineAnnealingLR

CosineAnnealingLR

如果只是了解是什么样的，很简单，一个公式：

$new_{lr}$表示新的学习率，$initial_{lr}$表示初始学习率，$eta_{min}$表示最小学习率，$T_{max}$表示 $cos$ 周期的$1/2$。

CosineAnnealingWarmRestarts

公式和CosineAnnealingLR是一样的，但是会在学习率达到最低点时restart，而且每次 restart 时会倍增周期。

import torch
from torchvision.models import AlexNet
from torch.optim.lr_scheduler import CosineAnnealingLR, CosineAnnealingWarmRestarts
import matplotlib.pyplot as plt
from lrs_scheduler import WarmRestart, warm_restart

model = AlexNet(num_classes=2)
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
# scheduler = CosineAnnealingLR(optimizer,T_max=20)
scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2)
# scheduler = WarmRestart(optimizer)
plt.figure()
x = list(range(100))
y = []
for epoch in range(1, 101):
    optimizer.zero_grad()
    optimizer.step()
    print("第%d个epoch的学习率：%f" % (epoch, optimizer.param_groups[0]['lr']))
    scheduler.step()
    # warm_restart(scheduler)
    y.append(scheduler.get_lr()[0])

# 画出lr的变化
plt.plot(x, y)
plt.xlabel("epoch")
plt.ylabel("lr")
plt.title("learning rate's curve changes as epoch goes on!")
plt.show()