plot_scheduling.py

import string
from copy import deepcopy

import torch
from torch import optim
import matplotlib.pyplot as plt

def main(schedulers,optimizerO,learning_rate,num_epochs,schedulerStep,title=""):

    plt.figure(figsize=(12, 5))
    for schedulerToUse in schedulers:
        optimizer=deepcopy(optimizerO)
        if schedulerToUse == "StepLR":
            scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=schedulerStep, gamma=0.75)
        elif schedulerToUse == "CosineAnnealing":
            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
                optimizer,
                T_max=num_epochs,
                eta_min=1e-6
            )
        elif schedulerToUse == "warmup+CosineAnnealing":
            warmup_epochs = 10
            min_lr = 1e-6

            warmup_scheduler = torch.optim.lr_scheduler.LinearLR(
                optimizer,
                start_factor=min_lr / learning_rate,
                end_factor=1.0,
                total_iters=warmup_epochs
            )

            cosine_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
                optimizer,
                T_max=num_epochs - warmup_epochs,
                eta_min=min_lr
            )

            scheduler = torch.optim.lr_scheduler.SequentialLR(
                optimizer,
                schedulers=[warmup_scheduler, cosine_scheduler],
                milestones=[warmup_epochs]
            )
        else:
            raise NotImplementedError
        x_epoch=[]
        y_epoch=[]

        for epoch in range(1,num_epochs+1):
            x_epoch.append(epoch)
            y_epoch.append(scheduler.get_last_lr()[0])
            print("lr :",y_epoch[-1], " at epoch : ",x_epoch[-1])

            optimizer.step()
            scheduler.step()
        plt.plot(x_epoch,y_epoch,label=schedulerToUse)
    plt.legend()
    plt.title(title)
    plt.grid(True)
    plt.show()

if __name__ == '__main__':
    policy=["CosineAnnealing","StepLR","warmup+CosineAnnealing"]
    lr=1e-3
    model=torch.nn.Sequential(torch.nn.Linear(10, 10), torch.nn.ReLU(), torch.nn.Linear(10, 10))
    optimizer = optim.Adam(model.parameters(), lr=lr)
    main(policy,optimizer,num_epochs=100,learning_rate=lr,schedulerStep=10,title=f"Used lr by epochs (initial lr= {lr})")