PyTorch_CustomExponentialLR

extension of pytorch exponential learning rate scheduler

• can use learning rate warmup by setting step_size_down > 0
• can use periodic learning rate schedule by setting step_size_up and step_size_down
• can use decaying periodic learning rate schedule by setting `max_lr_decay' < 1.

Usage

from custom_scheduler import CustomExponentialLR

optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = CustomExponentialLR(optimizer, base_lr=0.0001, max_lr=0.001, max_lr_decay=0.9, step_size_down=40, step_size_up=10)

for epoch in epochs:
    for x, target in dataloader:
        pred = model(x)
        loss = loss_f(pred, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    scheduler.step()

Visualize

import matplotlib.pyplot as plt

def visualize_scheduler(optimizer, scheduler, epochs):
    lrs = []
    for _ in range(epochs):
        optimizer.step()
        lrs.append(optimizer.param_groups[0]['lr'])
        scheduler.step()

    plt.plot(lrs)
    plt.show()
    
    return lrs
    

1. Exponential decay with warmup

from custom_scheduler import CustomExponentialLR
    
import torch
import torch.optim as optim

epochs = 100
optimizer = optim.SGD([torch.tensor(1)], lr=0.001, momentum=0.9)
scheduler = CustomExponentialLR(optimizer, base_lr=0.0001, max_lr=0.001, step_size_down=90, step_size_up=10)

lrs = visualize_scheduler(optimizer, scheduler, epochs)

2. Periodic exponential decay with warmup

from custom_scheduler import CustomExponentialLR
    
import torch
import torch.optim as optim

epochs = 100
optimizer = optim.SGD([torch.tensor(1)], lr=0.001, momentum=0.9)
scheduler = CustomExponentialLR(optimizer, base_lr=0.0001, max_lr=0.001, step_size_down=40, step_size_up=10)

lrs = visualize_scheduler(optimizer, scheduler, epochs)

3. Periodic, decaying exponential decay with warmup

from custom_scheduler import CustomExponentialLR
    
import torch
import torch.optim as optim

epochs = 100
optimizer = optim.SGD([torch.tensor(1)], lr=0.001, momentum=0.9)
scheduler = CustomExponentialLR(optimizer, base_lr=0.0001, max_lr=0.001, max_lr_decay=0.8, step_size_down=40, step_size_up=10)

lrs = visualize_scheduler(optimizer, scheduler, epochs)