使用带有不同LR的参数组的LR调度器的PyTorch

Question

我为每个参数组定义了以下具有不同学习率的优化器：

  optimizer = optim.SGD([
          {'params': param_groups[0], 'lr': CFG.lr, 'weight_decay': CFG.weight_decay},
          {'params': param_groups[1], 'lr': 2*CFG.lr, 'weight_decay': 0},
          {'params': param_groups[2], 'lr': 10*CFG.lr, 'weight_decay': CFG.weight_decay},
          {'params': param_groups[3], 'lr': 20*CFG.lr, 'weight_decay': 0},
      ], lr=CFG.lr, momentum=0.9, weight_decay=CFG.weight_decay, nesterov=CFG.nesterov)

现在，我想使用LR-Scheduler来更新所有学习率，而不仅仅是第一个学习率，因为默认情况下，调度器只会更新param_groups

scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=5, T_mult=2, eta_min=CFG.min_lr, last_epoch=-1, verbose=True)

给我：

Parameter Group 0
    dampening: 0
    initial_lr: 0.001
    lr: 0.0009999603905218616
    momentum: 0.9
    nesterov: True
    weight_decay: 0.0001

Parameter Group 1
    dampening: 0
    initial_lr: 0.002
    lr: 0.002
    momentum: 0.9
    nesterov: True
    weight_decay: 0

Parameter Group 2
    dampening: 0
    initial_lr: 0.01
    lr: 0.01
    momentum: 0.9
    nesterov: True
    weight_decay: 0.0001

Parameter Group 3
    dampening: 0
    initial_lr: 0.02
    lr: 0.02
    momentum: 0.9
    nesterov: True
    weight_decay: 0
)

在一次更新之后。

你知道怎么用调度程序更新所有的学习率吗？

Answer 1

你说得对，学习率调度程序应该逐个更新每个组的学习率。经过一些测试，看起来这个问题只出现在CosineAnnealingWarmRestarts调度器上。我测试了CosineAnnealingLR和其他几个调度器，它们更新了每个组的学习率：

 scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 100, verbose=True)

然后，为了找到问题的原因，我查看了学习速率调度器的源代码：https://github.com/pytorch/pytorch/blob/master/torch/optim/lr_scheduler.py

快速浏览一下，看起来在CosineAnnealingLR和CosineAnnealingWarmRestarts get_lr()函数的实现之间有一些不同：

 # CosineAnnealingLR:
 def get_lr(self):
     if not self._get_lr_called_within_step:
         warnings.warn("To get the last learning rate computed by the scheduler, "
                       "please use `get_last_lr()`.", UserWarning)
 
     if self.last_epoch == 0:
         return self.base_lrs
     elif (self.last_epoch - 1 - self.T_max) % (2 * self.T_max) == 0:
         return [group['lr'] + (base_lr - self.eta_min) *
                 (1 - math.cos(math.pi / self.T_max)) / 2
                 for base_lr, group in
                 zip(self.base_lrs, self.optimizer.param_groups)]
     return [(1 + math.cos(math.pi * self.last_epoch / self.T_max)) /
             (1 + math.cos(math.pi * (self.last_epoch - 1) / self.T_max)) *
             (group['lr'] - self.eta_min) + self.eta_min
             for group in self.optimizer.param_groups]    
 
 # CosineAnnealingWarmRestarts:
 def get_lr(self):
     if not self._get_lr_called_within_step:
         warnings.warn("To get the last learning rate computed by the scheduler, "
                       "please use `get_last_lr()`.", UserWarning)
 
     return [self.eta_min + (base_lr - self.eta_min) * (1 + math.cos(math.pi * self.T_cur / self.T_i)) / 2
             for base_lr in self.base_lrs]

所以看完代码后，我觉得这个问题是个bug。甚至CosineAnnealingWarmRestart的文档也建议“使用余弦退火时间表设置每个参数组的学习率”。