您可以使用 torch.optim.lr_scheduler 模块中的 LambdaLR 类来实现这个功能。
具体来说，您可以使用以下代码设置一个学习率调度器：
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# Define a function that returns the learning rate at each step
def lr_lambda(step):
if step < 2000:
return 1e-4 + (1e-3 - 1e-4) * step / 2000
else:
return 1e-3 - (1e-3 - 1e-4) * (step - 2000) / 8000
# Create a scheduler that applies the above function to the optimizer
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda)
# Train the model in a loop
for step in range(total_steps):
# ...
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
在这个示例中，我们首先使用 optim.Adam 创建一个优化器对象，并将学习率设置为 1e-3。
然后，我们定义一个函数 lr_lambda，它接受当前步骤数作为输入，并返回相应的学习率。
在前 2000 个步骤中，我们逐步将学习率从 1e-4 增加到 1e-3。在 2000 步之后，我们逐步将
学习率从 1e-3 降低到 1e-4。最后，我们使用 LambdaLR 创建一个学习率调度器，并在每个
训练步骤之后调用 scheduler.step()来更新优化器的学习率。
这样做的结果是，前 2000 步的学习率会很小，从而让模型慢慢地学习数据集的结构。随着
训练的进行，学习率会逐步增加，使得模型可以更快地收敛。
资源评论
