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

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是优化器的参数组，对于大多数情况，只有一个参数组。