在Pytorch中训练时重新设置学习率

Question

我正在Pytorch中使用学习率调度器训练一个模型，以降低学习率的值。通过使用学习率调度器，我将学习率从0.0001降低到1e-5，并在特定检查点保存所有权重、参数、学习率值等。现在，我想继续训练模型，但是学习率的值不同，其他值都保持不变。我该怎么做呢?这是保存检查点的代码。我使用Adam优化器

checkpoint = {
'epoch': epoch + 1,
'val_loss_min': val_loss['total'].avg,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
}

加载检查点时，我使用以下代码:

checkpoint = torch.load(args.SAVED_MODEL)
# Load current epoch from checkpoint
epochs = checkpoint['epoch']
# Load state_dict from checkpoint to model
model.load_state_dict(checkpoint['state_dict'])
# Load optimizer from checkpoint to optimizer
optimizer.load_state_dict(checkpoint['optimizer'])
# Load valid_loss_min from checkpoint to valid_loss_min
val_loss_min = checkpoint['val_loss_min']
# Load scheduler from checkpoint to scheduler
scheduler.load_state_dict(checkpoint['scheduler'])

Answer 1

您可以通过访问其param_groups属性来更改优化器的学习率。根据您是否有多个组，您可以执行以下操作(在将检查点加载到之后):

for g in optimizer.param_groups:
g['lr'] = new_lr