PyTorch学习率 warmup + 余弦退火

PyTorch学习率 warmup + 余弦退火

Pytorch 余弦退火

PyTorch内置了很多学习率策略，详情请参考torch.optim — PyTorch 1.10.1 documentation，这里只介绍常用的余弦退火学习率策略。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-noT4RJvn-1641352869753)(解决的问题记录-2021.assets/image-20220105105717934.png)]

主要参数：

optimizer：优化器，比如SGD、Adam之类的。

T_max：关键参数，余弦退火最大的迭代次数，即：学习率衰减到最小值时，迭代的次数，通常是len(dataloader)*epochs。

eta_min：退火过程中最小的学习率。

PyTorch warmup

Pytorch内部并没有warmup的接口，为此需要使用第三方包pytorch_warmup ，可以使用命令pip install pytorch_warmup进行安装。

改包支持多种学习策略：

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-mCGOpscU-1641352869754)(解决的问题记录-2021.assets/learning_rate.png)]

可以使用该软件包与Pytorch内置的学习率策略来完成warmup和余弦退火。

使用示例

官方给了两种使用warmup的方法：

1. 学习率策略使用全局迭代次数

   import torch
   import pytorch_warmup as warmup
   
   optimizer = torch.optim.AdamW(params, lr=0.001, betas=(0.9, 0.999), weight_decay=0.01)
   num_steps = len(dataloader) * num_epochs
   lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_steps)
   warmup_scheduler = warmup.UntunedLinearWarmup(optimizer)
   for epoch in range(1,num_epochs+1):
       for batch in dataloader:
           optimizer.zero_grad()
           loss = ...
           loss.backward()
           optimizer.step()
           lr_scheduler.step()
           warmup_scheduler.dampen()
   

对于PyTorch 大于1.4的们需要进行以下修改：

        lr_scheduler.step(lr_scheduler.last_epoch+1)

2. 学习率策略使用epoch数目（PyTorch >=1.2）

   lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[num_epochs//3], gamma=0.1)
   warmup_scheduler = warmup.UntunedLinearWarmup(optimizer)
   warmup_scheduler.last_step = -1 # initialize the step counter
   for epoch in range(1,num_epochs+1):
       for batch in dataloader:
           lr_scheduler.step(epoch-1)
           warmup_scheduler.dampen()
           optimizer.zero_grad()
           loss = ...
           loss.backward()
           optimizer.step()