PyTorch 中如何使用 CosineAnnealingLR？

推荐答案

在 PyTorch 中，CosineAnnealingLR 是一种学习率调度器，它根据余弦函数调整学习率。以下是如何使用 CosineAnnealingLR 的示例代码：

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

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

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

# 定义 CosineAnnealingLR 调度器
scheduler = CosineAnnealingLR(optimizer, T_max=50, eta_min=0)

# 训练循环
for epoch in range(100):
    # 训练步骤
    optimizer.zero_grad()
    output = model(torch.randn(10))
    loss = output.mean()
    loss.backward()
    optimizer.step()
    
    # 更新学习率
    scheduler.step()
    
    # 打印当前学习率
    print(f'Epoch {epoch+1}, Learning Rate: {scheduler.get_last_lr()[0]}')

本题详细解读

1. CosineAnnealingLR 的作用

CosineAnnealingLR 是一种学习率调度器，它根据余弦函数周期性地调整学习率。学习率会在每个周期内从初始值逐渐减小到最小值，然后再逐渐增大，如此循环往复。

2. 参数说明

• optimizer: 需要调整学习率的优化器。
• T_max: 一个周期的长度，即学习率从最大值到最小值再到最大值所经历的步数。
• eta_min: 学习率的最小值，默认为 0。

3. 使用步骤

1. 定义模型和优化器: 首先需要定义一个模型和一个优化器。
2. 创建调度器: 使用 CosineAnnealingLR 创建一个学习率调度器，并传入优化器、T_max 和 eta_min 参数。
3. 更新学习率: 在每个训练步骤或每个 epoch 结束时，调用 scheduler.step() 来更新学习率。

4. 示例代码解析

• 在示例代码中，我们定义了一个简单的线性模型和一个 SGD 优化器。
• 然后创建了一个 CosineAnnealingLR 调度器，设置 T_max=50 和 eta_min=0。
• 在训练循环中，每个 epoch 结束后调用 scheduler.step() 来更新学习率，并打印当前的学习率。

通过这种方式，学习率会根据余弦函数周期性地变化，从而帮助模型在训练过程中更好地收敛。