pytorch实现学习率衰减的方法总结

前言：

之前在训练模型时一直使用自定义的学习率衰减函数，如下所示：

def adjust_learning_rate(args, optimizer, epoch, gamma=0.1):
    """Sets the learning rate to the initial LR decayed 0.1 every 50 epochs"""
    lr = args.lr * (0.1 ** (epoch // 50))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

然后在每个epoch训练前调用一次，实现学习率的衰减：

    for epoch in range(args.start_epoch, args.epochs):
        adjust_learning_rate(args, optimizer, epoch, gamma=0.1)
        ......

后来发现pytorch也像tensorflow一样实现了很多学习率调整的方法，并封装在了torch.optim.lr_scheduler模块中，接下来对这些方法进行一下总结。

 

1. torch.optim.lr_scheduler.StepLR

* step_size (int): 学习率的衰减周期，单位是epoch。

* gamma(float): 学习率衰减的乘积因子，默认为0.1。

* 说明：等间隔的学习率衰减，每训练step_size个epoch，将学习率调整为lr=lr*gamma。

使用示例如下：

# Assuming optimizer uses lr = 0.05 for all groups
# lr = 0.05