torch.optim.lr_scheduler.StepLR用法

介绍

在深度学习模型训练过程中，学习率（learning rate）是一个非常重要的超参数，它决定了模型参数在每个迭代步骤中更新的幅度大小。通常情况下，我们会在训练过程中逐渐降低学习率，以使模型在接近最优解时更加稳定。torch.optim.lr_scheduler.StepLR 就是 PyTorch 中的一个学习率调度器，它可以在训练过程中按照一定的策略逐渐降低学习率。

使用示例

torch.optim.lr_scheduler.StepLR 的使用方法非常简单，我们只需要在创建优化器（optimizer）时将它作为参数传入，并指定降低学习率的策略。具体来说，StepLR 的构造函数有三个参数：

• optimizer：一个 PyTorch 优化器对象，如 torch.optim.SGD、torch.optim.Adam 等。
• step_size：降低学习率的间隔步数，即经过多少个迭代步骤后降低学习率。
• gamma：学习率降低的倍数，即学习率每次降低后的缩放比例。

下面是一个简单的示例：

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

model = ...
optimizer = optim.SGD(model.parameters(), lr=0.1)
scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

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

在这个示例中，我们首先定义了一个模型 model，并使用 SGD 优化器对模型进行训练，初始学习率为 0.1。然后，我们创建了一个 StepLR 学习率调度器，将它作为参数传入优化器。每个 epoch 结束后，我们调用 scheduler.step() 方法，以使学习率按照预先指定的策略进行调整。

在这个示例中，我们将 step_size 设置为 10，gamma 设置为 0.1，即每 10 个 epoch 降低一次学习率，降低的倍数为 0.1。这个策略可以根据实际情况进行调整。需要注意的是，StepLR 调度器只能对一个优化器进行学习率调整，如果需要对多个优化器进行调整，可以使用 torch.optim.lr_scheduler.MultiStepLR 或 torch.optim.lr_scheduler.ExponentialLR 等其他学习率调度器。