Gradual Warmup Scheduler

Gradual Warmup Scheduler

1、介绍

从最开始的小学习率开始，每个iteration增大一点，直到最初设置的比较大的学习率。

其学习率如上图变化所示。

2、实现

这将使用到torch.optim.lr_scheduler中的类，比如torch.optim.lr_scheduler.LambdaLR()

CLASS torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch=- 1, verbose=False)

它具有一下几个参数：

• optimizer ，需要调节的优化器
• lr_lambda ，lambda系数的获取函数或列表
• last_epoch ，上一个epoch的索引
• verbose ，如果为True ，则为每次更新打印一条消息

参数更新：

$new\_lr=\lambda \times initial\_lr$

而lr_lambda参数传入一个函数，或是列表。

例如

total_epoch = 5
total_step = total_epoch * len(mnist_data_loader_train)
scheduler = LambdaLR(optimizer=optimizer,
                     lr_lambda=lambda step: (step + 1) / total_step if step < total_step else 1)

例如计划在第5个epoch使得学习率达到给定的初始值，那么$\lambda$则可以从一个小数开始，增加至1即可。

则这里的

lambda step: (step + 1) / total_step if step < total_step else 1

就是
$$\{\begin{array}{c}\frac{step+1}{total\_step},step<total\_step\\ 1,step>total\_step\end{array}$$
而学习率的更新，可以在每一个step更新，也可以每一个epoch更新一次

每一epoch更新一次：

for epoch in range(100):
     train(...)
     validate(...)
     scheduler.step()

或是每一个step更新：

for epoch in range(100):
     '''train(...)'''
     for step in data_loader:
        train()
        scheduler.step()
     validate(...)

3、完整的一个训练mnist实例：

import torchvision.transforms
from torch import nn
import torch
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torch.optim.lr_scheduler import LambdaLR
from tqdm import tqdm


class Model(nn.Module):  # 定义网络

    def __init__(self):
        super(Model, self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(1, 32, 3, 2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, 2),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, 2),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(512, 10)
        )

    def forward(self, x):
        return self.layer(x)


if __name__ == '__main__':
    """参数"""
    batch_size = 256
    lr = 0.01  # 学习率
    epochs = 100  # 迭代次数
    total_epoch = 5  # 总共经过total_epoch到达最终的学习率

    model = Model()  # 实例化模型
    mnist_train = MNIST('data', download=True, train=True, transform=torchvision.transforms.ToTensor())  # 训练集
    mnist_val = MNIST('data', download=True, train=False, transform=torchvision.transforms.ToTensor())  # 验证集
    mnist_data_loader_train = DataLoader(mnist_train, batch_size, True, drop_last=True)  # 训练集加载器
    mnist_data_loader_val = DataLoader(mnist_val, batch_size, True, drop_last=False)  # 验证集加载器

    optimizer = torch.optim.Adam(model.parameters(), lr)  # 优化器
    total_step = total_epoch * len(mnist_data_loader_train)
    scheduler = LambdaLR(optimizer=optimizer,
                         lr_lambda=lambda step: (step + 1) / total_step if step < total_step else 1)  # 学习速率调度器
    loss_function = nn.CrossEntropyLoss(label_smoothing=0.1)  # 损失函数
    accuracy = 0
    with tqdm(range(epochs), dynamic_ncols=True) as tqdm_range:
        for epoch in tqdm_range:  # 训练迭代过程
            loss_sum = 0
            for i, (x, y_hat) in enumerate(mnist_data_loader_train):
                y = model(x)  # 模型预测
                loss = loss_function(y, y_hat)  # 预测值与真实值做损失

                optimizer.zero_grad()  # 清空梯度
                loss.backward()  # 方向传播
                optimizer.step()  # 梯度下降

                loss_sum += loss.detach().item()  # 统计损失
                scheduler.step()  # 调整学习速率，也可以写在该层for循环外，根据epoch改变学习率
                tqdm_range.set_postfix(
                    ordered_dict={
                        'Epoch': epoch,
                        'Batch': f'{i + 1}/{len(mnist_data_loader_train)}',
                        'Lr': optimizer.param_groups[0]['lr'],
                        'Loss': loss_sum / (i + 1),
                        'Accuracy': accuracy
                    }
                )

            """验证"""
            correct_sum = 0
            for i, (x, y_hat) in enumerate(mnist_data_loader_val):
                y = model(x)
                y = torch.argmax(y, dim=1)
                correct_sum += (y == y_hat).sum().item()
            accuracy = f"{str(correct_sum / len(mnist_val) * 100)[:6]}%"
            print()
y, dim=1)
                correct_sum += (y == y_hat).sum().item()
            accuracy = f"{str(correct_sum / len(mnist_val) * 100)[:6]}%"
            print()