torch.optim.lr_scheduler.StepLR 的参数详解和应用

2023-12-30 01:08:07

torch.optim.lr_scheduler.StepLR 是 PyTorch 中学习率调度器(learning rate scheduler)的一种,它按照指定的步骤降低学习率。以下是 StepLR 的主要参数解释:

torch.optim.lr_scheduler.StepLR(optimizer, step_size, gamma=0.1, last_epoch=-1, verbose=False)
  • optimizer (torch.optim.Optimizer):要进行学习率调整的优化器。

  • step_size (int):调整学习率的步长,即经过多少个epoch后,学习率进行一次调整。

  • gamma (float, 默认值为 0.1):每次学习率调整的乘数因子。新的学习率 = 旧的学习率 * gamma。

  • last_epoch (int, 默认值为 -1):上一个epoch的索引,用于指定从哪个epoch开始学习率调整。如果设置为 -1,表示从头开始。

  • verbose (bool, 默认值为 False):如果设置为 True,将会在每次学习率调整时打印一条消息。

在 PyTorch 中,你可以通过将学习率调度器与优化器一起使用,实现在训练过程中动态调整学习率。以下是在模型中使用学习率调度器的一般步骤:

  1. 定义模型: 首先,定义你的神经网络模型。

  2. 定义优化器: 创建一个优化器(如 SGD 或 Adam)并将模型的参数传递给它。

  3. 定义学习率调度器: 创建一个学习率调度器,并将其与优化器关联。

  4. 在训练循环中使用: 在每个训练迭代(或每个 epoch)结束时,调用学习率调度器的 step()方法。

下面是一个简单的示例:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR

# 步骤1: 定义模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 2)

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

# 步骤2: 定义优化器
model = SimpleModel()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 步骤3: 定义学习率调度器
scheduler = StepLR(optimizer, step_size=5, gamma=0.1)

# 步骤4: 在训练循环中使用
num_epochs = 20

for epoch in range(num_epochs):
    # 模型训练步骤
    # ...

    # 在每个epoch结束时调整学习率
    scheduler.step()

    # 打印当前学习率
    print(f"Epoch {epoch+1}/{num_epochs}, Learning Rate: {optimizer.param_groups[0]['lr']:.5f}")

在上述示例中,StepLR 调度器在每5个 epoch 后将学习率乘以0.1。在每个 epoch 结束时,调用 scheduler.step()来更新学习率。在训练循环中,你可以通过optimizer.param_groups[0]['lr']获取当前的学习率。这里的 param_groups是优化器的参数组,对于大多数情况,只有一个参数组。

文章来源:https://blog.csdn.net/weixin_51659315/article/details/135203341
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。