cyclic learning rate with multiple learning rates

Viewed 308

I am trying to use the OneCycleLR or atleast the cyclicLR in torch.optim.lr_scheduler.

Suppose I have the following:

param_list = []
for lr, block in zip(lrs, blocks):
    param_list.extend([{'params':p ,'lr':lr} for n,p in model.named_parameters() if n.startswith(block)])
optimizer = torch.optim.Adam(param_list)

where blocks = ["base", "fc"] (in my use case there is ~20 blocks) and lrs=[1e-4, 1e-3].

It is easy enough to control the learning rates manually by using a function, eg:

lr_sched = lambda batch: 1.1**batch
scheduler = LambdaLR(optimizer, lr_lambda=[lr_sched]*len(param_list))

The above example increases the learning rate.

However, what I would like to do is change the learning rate as well as the momentum parameters as offered in OneCycleLR. So my question is:

  1. Is it possible?
  2. If not, is there a way to manipulate the momentum while training, and I can write a function for the cyclic learning rate myself.
  3. Is it possible to use a list of optimizers instead of one, if so is that slower?

Minimal example:

import torch
import torch.nn as nn
from torch.optim.lr_scheduler import LambdaLR

class Model(nn.Module):
    def __init__(self):
        self.base = nn.Linear(10, 5)
        self.fc = nn.Linear(5, 1)
        self.relu = nn.ReLU()
        
    def forward(self, x):
        return self.fc(self.relu(self.base(x)))

model = Model()

param_list = []
for lr, block in zip(lrs, blocks):
    param_list.extend([{'params':p ,'lr':lr} for n,p in model.named_parameters() if n.startswith(block)])
optimizer = torch.optim.Adam(param_list)

lr_sched = lambda batch: 1.1**batch
scheduler = LambdaLR(optimizer, lr_lambda=[lr_sched]*len(param_list))
0 Answers
Related