PyTorch 使用具有不同 LR 参数组的 LR-Scheduler

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[0]？

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]

所以看了代码之后，我觉得这个问题是一个错误。甚至 CosineAnnealingWarmRestart 的文档也建议“使用余弦退火计划设置每个参数组的学习率”。