模型训练技巧——warm up


1. pytorch 中学习率的调节策略

（1）等间隔调整学习率 StepLR

（2）按需调整学习率 MultiStepLR

（3）指数衰减调整学习率 ExponentialLR

（4）余弦退火调整学习率 CosineAnnealingLR

（5）自适应调整学习率 ReduceLROnPlateau

（6）自定义调整学习率 LambdaLR

每种学习率的参数详解，见博文：pytorch 学习率参数详解

2. 论文中和比赛中学习率的调节策略

然而在顶会论文和知名比赛中，作者一般都不会直接使用上述学习率调整策略，而是先预热模型（warm up）, 即以一个很小的学习率逐步上升到设定的学习率，这样做会使模型的最终收敛效果更好。

下面，小编以warm up + CosineAnnealingLR来实现学习率的调整。训练过程中学习率的变化过程如图中红色曲线所示：

Caption

3. 代码实现

首先，写一个warm up的类，重写get_lr方法。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import torch
from torch.optim.lr_scheduler import _LRScheduler


class WarmUpLR(_LRScheduler):
    """warmup_training learning rate scheduler

    Args:
        optimizer: optimzier(e.g. SGD)
        total_iters: totoal_iters of warmup phase
    """
    def __init__(self, optimizer, total_iters, last_epoch=-1):
       
        self.total_iters = total_iters
        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        """we will use the first m batches, and set the learning
        rate to base_lr * m / total_iters
        """
        return [base_lr * self.last_epoch / (self.total_iters + 1e-8) for base_lr in self.base_lrs]

在训练代码中使用：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9, weight_decay=5e-4)

    warmup_epoch = 5
    scheduler = CosineAnnealingLR(optimizer, 100 - warmup_scheduler)
   
    iter_per_epoch = len(train_dataset)
    warmup_scheduler = WarmUpLR(optimizer, iter_per_epoch * warmup_epoch)

    for epoch in range(1, max_epoch+1):
        if epoch >= warmup_epoch:
            scheduler.step()
            learn_rate = scheduler.get_lr()[0]
            print("Learn_rate:%s" % learn_rate)
        test(epoch, net, valloader, criterion)
        train(epoch, net, trainloader, optimizer, criterion, warmup_scheduler)

在train函数中的修改：

1
2
3
4
5
6
for (inputs, targets) in tqdm(trainloader):
        if epoch < 5:
            warmup_scheduler.step()
            warm_lr = warmup_scheduler.get_lr()
            print("warm_lr:%s" % warm_lr)
        inputs, targets = inputs.to(device), targets.to(device)

4. 总结

在论文中和比赛中一般都会用到warm up技巧，特别是在模型难收敛的任务中。在论文中，MultiStepLR和CosineAnnealingLR两种学习率调节策略用得较多。在知名竞赛中，ReduceLROnPlateau学习率调整策略用得较多。小编在工程项目中是怎么用的呢？一般用warm up结合上述三种调节策略都尝试一遍，最终哪个模型的精度高就用哪个模型。很多情况下，三个模型的精度差不多，精度差距在±0.5%以内。