cyclic learning rate with multiple learning rates - pytorch

Question

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:
Is it possible?
If not, is there a way to manipulate the momentum while training, and I can write a function for the cyclic learning rate myself.
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))