torch.optim.lr_scheduler.StepLR怎么使用

torch.optim.lr_scheduler.StepLR 是 PyTorch 中的一个学习率调度器，它将学习率按照给定的步长下降。使用该调度器可以让模型在训练过程中逐渐降低学习率，以便更好地收敛到最优解。

下面是使用 torch.optim.lr_scheduler.StepLR 的示例代码：

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

# 定义优化器和学习率调度器
optimizer = optim.SGD(model.parameters(), lr=0.1)
scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

# 在训练过程中更新学习率
for epoch in range(num_epochs):
    # 训练模型
    train_loss = ...
    train_acc = ...
    
    # 更新学习率
    scheduler.step()
    
    # 在验证集上评估模型
    val_loss = ...
    val_acc = ...
    
    # 打印训练和验证结果
    print("Epoch [{}/{}], Train Loss: {:.4f}, Train Acc: {:.4f}, Val Loss: {:.4f}, Val Acc: {:.4f}"
          .format(epoch+1, num_epochs, train_loss, train_acc, val_loss, val_acc))

在上面的代码中，我们首先定义了一个 SGD 优化器，然后使用 StepLR 调度器将其包装起来。我们将步长设置为 10，即每 10 个 epoch 降低学习率。gamma 参数是一个降低因子，它将当前学习率乘以 gamma 来得到新的学习率。

在每个 epoch 结束时，我们调用 scheduler.step() 来更新学习率。这样，优化器就会使用新的学习率来更新模型参数。

需要注意的是，调用 scheduler.step() 的位置很重要。在上面的示例代码中，我们将其放在了每个 epoch 结束后。如果你将其放在每个 batch 结束后，那么学习率将会在每个 batch 中更新，这可能会导致模型难以收敛。