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

如何在PyTorch中使用torch.optim.lr_scheduler_LRScheduler()来调整学习率

发布时间:2023-12-29 15:02:46

在PyTorch中,torch.optim.lr_scheduler模块提供了多种学习率调度器,用于在训练过程中自动调整学习率。其中,torch.optim.lr_scheduler.LambdaLR是一个非常常用的调度器,它通过定义一个学习率调整函数来改变学习率。

首先,我们需要导入必要的库和模块:

import torch
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler

接下来,我们可以定义一个简单的神经网络模型,并定义一个损失函数和优化器。请注意,为了方便演示,我将使用一个单层全连接网络和MSE损失函数作为示例:

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = torch.nn.Linear(10, 1)

    def forward(self, x):
        out = self.fc(x)
        return out

net = Net()
loss_fn = torch.nn.MSELoss()
optimizer = optim.SGD(net.parameters(), lr=0.1)

接着,我们可以定义一个学习率调整函数。这个函数接受当前的学习率和当前的训练轮数作为输入,并返回一个新的学习率。

def adjust_learning_rate(optimizer, epoch):
    lr = 0.1  # 初始学习率

    # 在第50和第100轮时,将学习率降低为原来的1/10
    if epoch > 100:
        lr *= 0.1
    elif epoch > 50:
        lr *= 0.01

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

然后,我们可以使用上面定义的学习率调整函数在每个训练轮次中调整学习率。可以在训练循环的每个epoch中添加以下代码:

for epoch in range(200):
    adjust_learning_rate(optimizer, epoch)
    # 正常的训练代码
    for inputs, labels in dataloader:
        outputs = net(inputs)
        loss = loss_fn(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

在上述代码中,我们通过调用adjust_learning_rate函数来动态调整学习率。

除了LambdaLR,torch.optim.lr_scheduler模块还提供了其他学习率调度器,如StepLR、MultiStepLR、ExponentialLR和CosineAnnealingLR等。这些调度器通过不同的方式调整学习率,可以根据具体需求进行选择和使用。

希望通过上述示例,您能了解如何在PyTorch中使用torch.optim.lr_scheduler.LambdaLR来调整学习率。请记住,在实际应用中,根据不同的模型和任务,可能需要调整学习率的具体策略和参数。