How to merge two learning rate schedulers in PyTorch?

Viewed 2551

I want to use OneCycleLR for warming up the LR from a low value, then decreasing the LR from peak with ReduceLROnPlateau.

I understand that I could simulate ReduceLROnPlateau manually after increasing LR with OneCycleLR, but is there a simpler way to use these two schedulers together directly?

Thanks in advance.

1 Answers

PyToch has released a method, on github instead of official guidelines.

You can try the following snippet:

import torch
from torch.nn import Parameter
from torch.optim import SGD
from torch.optim.lr_scheduler import ExponentialLR, StepLR

model = [Parameter(torch.randn(2, 2, requires_grad=True))]
optimizer = SGD(model, 0.1)

scheduler1 = ExponentialLR(optimizer, gamma=0.9)
scheduler2 = StepLR(optimizer, step_size=5, gamma=0.1)

for epoch in range(10):

    print(epoch, scheduler2.get_last_lr()[0])

    optimizer.step()
    scheduler1.step()
    scheduler2.step()
Related