python - 每次重启后降低最大学习率

Question

我正在为基于计算机视觉的任务训练神经网络。对于优化器，我发现在整个训练中使用单一学习率并不理想，人们所做的是他们使用学习率调度程序以特定方式衰减学习率。所以为了做到这一点，我尝试了PyTorch's CosineAnnealingWarmRestarts().它的作用是它以余弦方式退火/降低初始学习率（由我们设置），直到它重新启动。在这个“重启”之后，学习率被设置回初始学习率，并且循环再次发生。这对我来说效果很好，但我想对其进行一些更改。我想改变学习率，每次重启后都会分配优化器，这样每次重启后优化器的最大学习率也会降低。这可以在 PyTorch 中完成吗？

Answer 1

在我看来，一个直截了当的解决方案就是继承CosineAnnealingWarmRestarts然后self.optimizer在覆盖step函数中更改其参数。在伪代码中，这将类似于

class myScheduler(torch.optim.lr_scheduler.CosineAnnealingWarmRestarts):
    def __init__(self,
            optimizer, 
            T_0, 
            T_mult=1,
            eta_min=0,
            last_epoch=-1):
    #initialize base calss
    super().__init__(.... blablabla ...)

    def step(self):
        #call step() from base class
        #Do some book-keeping to determine if you've hit a restart 
        #now change optimizer lr for each parameter group
        if some_condition:#condition like number of iterations, restarts, etc
            self.optimizer.param_groups[i]['lr']*=some_coef