torch.optim.lr_scheduler.StepLR()函数

1 目的

在训练的开始阶段, 使用的 LR 较大, loss 可以下降的较快, 但是随着训练的轮数越来越多, loss 越来越接近 global min, 若不降低 LR , 就会导致 loss 在最小值左右来回跳动(因为 LR 较大的缘故), 因此, 随着训练轮数的增加, 我们因该逐步减小 LR, 以更快的达到 global min 并减小局部震荡, 这就是 optim.lr_scheduler.StepLR() 函数的作用.

2 StepLR() 及其参数介绍

import torch.optim as optim
scheduler_lr = optim.lr_scheduler.StepLR(optimizer, step_size, gamma, last_epoch = -1, verbose = False)

# 记得在训练开始进行激活
optimizer.zero_grad()
optimizer.step()
scheduler_lr.step()

参数介绍:

• optimizer: 神经网络所使用的优化器

• step_size: 多少轮循环后更新一次学习率

• gamma: 每次将 lr 更新为原来的 gamma 倍