PyTorch 中如何使用 MultiStepLR？

推荐答案

在 PyTorch 中，MultiStepLR 是一种学习率调度器，它允许在训练过程中根据预定义的里程碑（milestones）来调整学习率。以下是使用 MultiStepLR 的示例代码：

// www.javascriptcn.com code example
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import MultiStepLR

# 定义一个简单的模型
model = nn.Linear(10, 1)

# 定义优化器
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 定义学习率调度器，设置里程碑为 [30, 80]，即在第30和第80个epoch时调整学习率
scheduler = MultiStepLR(optimizer, milestones=[30, 80], gamma=0.1)

# 训练循环
for epoch in range(100):
    # 训练步骤...
    
    # 更新学习率
    scheduler.step()
    
    # 打印当前学习率
    print(f'Epoch {epoch+1}, Learning Rate: {scheduler.get_last_lr()}')

本题详细解读

1. MultiStepLR 的作用

MultiStepLR 是一种学习率调度器，它允许在训练过程中根据预定义的里程碑（milestones）来调整学习率。当训练的 epoch 数达到里程碑时，学习率会乘以 gamma 参数指定的因子。

2. 参数说明

• optimizer: 需要调整学习率的优化器。
• milestones: 一个列表，指定在哪些 epoch 时调整学习率。例如，[30, 80] 表示在第30和第80个epoch时调整学习率。
• gamma: 学习率调整的因子。例如，gamma=0.1 表示学习率会乘以0.1。

3. 使用步骤

1. 定义模型和优化器。
2. 创建 MultiStepLR 调度器，并传入优化器、里程碑和 gamma 参数。
3. 在训练循环中，每个 epoch 结束后调用 scheduler.step() 来更新学习率。

4. 注意事项

• scheduler.step() 应该在每个 epoch 结束后调用，而不是在每个 batch 结束后调用。
• scheduler.get_last_lr() 可以获取当前的学习率，方便调试和记录。

通过使用 MultiStepLR，可以在训练过程中动态调整学习率，从而更好地优化模型。