欢迎访问宙启技术站
智能推送

Python中基于LRScheduler()的学习率调整策略

发布时间:2023-12-13 05:28:57

学习率在深度学习中起到至关重要的作用,它决定了每一次参数更新的幅度。如果学习率设置得太小,训练过程会变得非常缓慢;而如果学习率设置得太大,则可能导致模型无法收敛或者在最优解附近震荡。为了解决这个问题,在训练过程中,我们需要根据训练的进程动态地调整学习率。

PyTorch中提供了一个非常方便的学习率调整策略类LRScheduler(),该类可以与优化器(optimizer)配合使用,从而实现自动调整学习率的功能。LRScheduler()中包含了一些常用的学习率调整策略,例如StepLR()、MultiStepLR()、ExponentialLR()等。

下面我们来介绍几种常用的学习率调整策略及其使用方法。

1. StepLR()

StepLR是一种简单的学习率调整策略,它在训练过程中按照固定的步数(step_size)进行学习率的更新,更新方式为将学习率乘以gamma。

使用方法如下:

scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size, gamma)

其中,optimizer是我们定义的优化器,step_size表示学习率更新的步数,gamma表示更新时学习率的缩放因子。

2. MultiStepLR()

MultiStepLR是一种根据给定的milestones进行学习率调整的策略。具体来说,我们可以在训练过程中设定多个milestones,当当前epoch数达到某个milestone时,学习率将按照给定的gamma进行调整。

使用方法如下:

scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones, gamma)

其中,optimizer是我们定义的优化器,milestones是包含milestone数值的列表,gamma表示更新时学习率的缩放因子。

3. ExponentialLR()

ExponentialLR是一种根据指数函数进行学习率调整的策略。具体来说,学习率将按照给定的gamma指数衰减。

使用方法如下:

scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma)

其中,optimizer是我们定义的优化器,gamma表示更新时学习率的缩放因子。

下面给出一个使用例子,演示如何在训练过程中使用LRScheduler()进行学习率的动态调整。

import torch
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
import torchvision

# 定义一个简单的神经网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784, 64)
        self.fc2 = nn.Linear(64, 10)
        
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 加载MNIST数据集
transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)

# 初始化模型、优化器和损失函数
model = Net()
optimizer = optim.SGD(model.parameters(), lr=0.1)
criterion = nn.CrossEntropyLoss()

# 定义学习率调整策略,使用MultiStepLR,在第30和第60个epoch时将学习率分别缩小为原来的0.1和0.01倍
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[30, 60], gamma=0.1)

# 训练模型
for epoch in range(100):
    for i, (inputs, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(inputs.view(-1, 784))
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
    scheduler.step()
    print('Epoch: {}, Learning Rate: {}'.format(epoch, scheduler.get_lr()[0]))

在上述例子中,我们定义了一个简单的神经网络模型,并加载了MNIST数据集。然后我们使用SGD优化器作为我们的优化器,并使用CrossEntropyLoss作为损失函数。创建了一个学习率调整策略- MultiStepLR,在第30和第60个epoch时将学习率分别缩小为原来的0.1和0.01倍。在训练过程中,我们调用scheduler.step()来更新学习率,并使用scheduler.get_lr()来获取当前的学习率。

以上就是使用LRScheduler()进行学习率调整的一些常用策略及其使用方法。通过合理使用学习率调整策略,我们可以加快模型的训练速度,并提高模型的性能。