"""
 <h1 style="font-family:verdana;"> <center>BERT BASELINE AND A WALKTHROUGH THE LEARNING RATE SCHEDULERS </center> </h1>
"""
"""

<p style="color:#159364; font-family:cursive;">INSTALL THE TRANSFORMERS PACKAGE FROM THE HUGGING FACE LIBRARY</center></p>

"""
!pip install transformers
"""
# <p style="color:#159364; font-family:cursive;">IMPORT THE LIBRARIES</center></p>
"""
import os
import gc
import copy
import datetime
import time
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import Dataset, DataLoader
from torch.cuda import amp
import transformers
from transformers import BertTokenizer,BertForSequenceClassification, BertModel, BertConfig
from transformers import AdamW, get_linear_schedule_with_warmup
from tqdm import tqdm
from collections import defaultdict
import plotly.graph_objects as go
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import StratifiedKFold, KFold
import warnings
warnings.filterwarnings("ignore")

"""
# <p style="color:#159364; font-family:cursive;">LOOK AT THE DATA</center></p>
"""
df = pd.read_csv("../input/commonlitreadabilityprize/train.csv")
test_df = pd.read_csv("../input/commonlitreadabilityprize/test.csv",usecols=["id","excerpt"])
print('Number of training sentences: {:,}\n'.format(df.shape[0]))
df.sample(10)
"""
# <p style="color:#159364; font-family:cursive;">A BIT OF PREPROCESSING</center></p>
"""
def prep_text(text_df):
    text_df = text_df.str.replace("\n","",regex=False) 
    return text_df.str.replace("\'s",r"s",regex=True).values
df["excerpt"] = prep_text(df["excerpt"])
test_df["excerpt"] = prep_text(test_df["excerpt"])
"""
# <p style="color:#159364; font-family:cursive;">CREATE FOLDS</center></p>
"""
"""
Code taken from:https://www.kaggle.com/abhishek/step-1-create-folds
"""
def create_folds(data, num_splits):
    # we create a new column called kfold and fill it with -1
    data["kfold"] = -1
    
    # the next step is to randomize the rows of the data
    data = data.sample(frac=1).reset_index(drop=True)

    # calculate number of bins by Sturge's rule
    # I take the floor of the value, you can also
    # just round it
    num_bins = int(np.floor(1 + np.log2(len(data))))
    
    # bin targets
    data.loc[:, "bins"] = pd.cut(
        data["target"], bins=num_bins, labels=False
    )
    
    # initiate the kfold class from model_selection module
    kf = StratifiedKFold(n_splits=num_splits)
    
    # fill the new kfold column
    # note that, instead of targets, we use bins!
    for f, (t_, v_) in enumerate(kf.split(X=data, y=data.bins.values)):
        data.loc[v_, 'kfold'] = f
    
    # drop the bins column
    data = data.drop("bins", axis=1)

    # return dataframe with folds
    return data


# create folds
df = create_folds(df, num_splits=5)
"""
# <p style="color:#159364; font-family:cursive;">TRAINING CONFIGURATION</center></p>
"""
sen_length = []

for sentence in tqdm(df["excerpt"]):
   
    token_words = CONFIG.tokenizer.encode_plus(sentence)["input_ids"]
    sen_length.append(len(token_words))
    
print('maxlenth of all sentences are  ', max(sen_length))
class CONFIG:
    seed = 42
    max_len = 331
    train_batch = 16
    valid_batch = 32
    epochs = 10
    learning_rate = 2e-5
    splits = 5
    scaler = amp.GradScaler()
    model='bert-base-cased'
    tokenizer = BertTokenizer.from_pretrained(model, do_lower_case=True)
    tokenizer.save_pretrained('./tokenizer')
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
"""
# <p style="color:#159364; font-family:cursive;">REPRODUCIBILITY</center></p>
"""
def set_seed(seed = CONFIG.seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    os.environ['PYTHONHASHSEED'] = str(seed)
    
set_seed(CONFIG.seed)
"""
# <p style="color:#159364; font-family:cursive;">DEFINE THE DATASET CLASS</center></p>
"""
class BERTDataset(Dataset):
    def __init__(self,df):
        self.text = df['excerpt'].values
        self.target = df['target'].values
        self.max_len = CONFIG.max_len
        self.tokenizer = CONFIG.tokenizer
        
    def __len__(self):
        return len(self.text)
    
    def __getitem__(self, index):
        text = self.text[index]
        text = ' '.join(text.split())
        inputs = self.tokenizer.encode_plus(
            text,
            None,
            truncation=True,
            add_special_tokens=True,
            max_length=self.max_len,
            padding='max_length',
            return_token_type_ids=True
        )

        return {
            'ids': torch.tensor(inputs['input_ids'], dtype=torch.long),
            'mask': torch.tensor(inputs['attention_mask'], dtype=torch.long),
            'token_type_ids': torch.tensor(inputs["token_type_ids"], dtype=torch.long),
            'target': torch.tensor(self.target[index], dtype=torch.float)
        }
"""
# <p style="color:#159364; font-family:cursive;">MODEL:BERT FOR SEQUENCE CLASSIFICATION from 🤗 </center></p>
"""
model = BertForSequenceClassification.from_pretrained(
    "bert-base-uncased",
    num_labels = 1,
    output_attentions = False,
    output_hidden_states = False, 
)

model.cuda()
"""
# <p style="color:#159364; font-family:cursive;">OPTIMIZER</center></p>
"""
param_optimizer = list(model.named_parameters())
no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
optimizer_parameters = [
    {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 
     'weight_decay': 0.0001},
    {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 
     'weight_decay': 0.0}
    ]  

optimizer = AdamW(optimizer_parameters, lr=CONFIG.learning_rate)
"""
# <p style="color:#159364; font-family:cursive;">GET THE PREPARED DATA</center></p>
"""
def get_data(fold):
    df_train = df[df.kfold != fold].reset_index(drop=True)
    df_valid = df[df.kfold == fold].reset_index(drop=True)
    
    train_dataset = BERTDataset(df_train)
    valid_dataset = BERTDataset(df_valid)

    train_loader = DataLoader(train_dataset, batch_size=CONFIG.train_batch, 
                              num_workers=4, shuffle=True, pin_memory=True)
    valid_loader = DataLoader(valid_dataset, batch_size=CONFIG.valid_batch, 
                              num_workers=4, shuffle=False, pin_memory=True)
    
    return train_loader, valid_loader
"""
# <p style="color:#159364; font-family:cursive;">FOLD:0</center></p>
"""
train_dataloader,validation_dataloader=get_data(0)
"""
# <p style="color:#159364; font-family:cursive;">DEFINE LOSS AND TIME FUNCTIONS</center></p>
"""
def loss_fn(output,target):
     return torch.sqrt(nn.MSELoss()(output,target))
def format_time(elapsed):
    elapsed_rounded = int(round((elapsed)))
    return str(datetime.timedelta(seconds=elapsed_rounded))
"""
# <p style="color:#159364; font-family:cursive;">DEFINE THE FUNCTION FOR TRAINING,VALIDATION AND RUNNING</center></p>
"""
def run(optimizer,scheduler):
    set_seed(40)
    scaler=CONFIG.scaler
    training_stats = []
    total_t0 = time.time()
    epochs=CONFIG.epochs
    for epoch_i in range(0, epochs):
        print("")
        print('======== Epoch {:} / {:} ========'.format(epoch_i + 1, epochs))
        print('Training...')
        t0 = time.time()
        total_train_loss = 0
        data_size=0
        model.train()
        for step, batch in enumerate(train_dataloader):
            tr_loss=[]
            b_input_ids = batch['ids'].to(CONFIG.device)
            b_token_type_ids=batch['token_type_ids'].to(CONFIG.device)
            b_input_mask = batch['mask'].to(CONFIG.device)
            b_labels = batch['target'].to(CONFIG.device)
            batch_size = b_input_ids.size(0)
            model.zero_grad() 
            with amp.autocast(enabled=True):
                output= model(b_input_ids,attention_mask=b_input_mask)          
                output=output["logits"].squeeze(-1)
                loss = loss_fn(output,b_labels)
                tr_loss.append(loss.item()/len(output))
            scheduler.step()
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

        avg_train_loss = np.mean(tr_loss)    
        training_time = format_time(time.time() - t0)
        gc.collect()
        print("")
        print("  Average training loss: {0:.2f}".format(avg_train_loss))
        print("  Training epoch took: {:}".format(training_time))
        print("")
        print("Running Validation...")

        t0 = time.time()
        model.eval()
        val_loss = 0
        allpreds = []
        alltargets = []
        for batch in validation_dataloader:
            losses = []
            with torch.no_grad():
                device=CONFIG.device
                ids = batch["ids"].to(device)
                mask = batch["mask"].to(device)
                tokentype = batch["token_type_ids"].to(device)
                output = model(ids,mask)
                output = output["logits"].squeeze(-1)
                target = batch["target"].to(device)
                loss = loss_fn(output,target)
                losses.append(loss.item()/len(output))
                allpreds.append(output.detach().cpu().numpy())
                alltargets.append(target.detach().squeeze(-1).cpu().numpy())
        allpreds = np.concatenate(allpreds)
        alltargets = np.concatenate(alltargets)
        losses = np.mean(losses)
        gc.collect() 
        validation_time = format_time(time.time() - t0)
        print("  Validation Loss: {0:.2f}".format(losses))
        print("  Validation took: {:}".format(validation_time))
        training_stats.append(
        {
            'epoch': epoch_i + 1,
            'Training Loss': avg_train_loss,
            'Valid. Loss': losses,
            'Training Time': training_time,
            'Validation Time': validation_time
        }
    ) 
    print("")
    print("Training complete!")
    return training_stats  
"""
# <p style="color:#159364; font-family:cursive;">VISUALIZATION FUNCTION </center></p>
"""
def Visualizations(training_stats):
    pd.set_option('precision', 2)
    df_stats = pd.DataFrame(data=training_stats)
    df_stats = df_stats.set_index('epoch')
    layout = go.Layout(template= "plotly_dark")
    fig = go.Figure(layout=layout)
    fig.add_trace(go.Scatter(x=df_stats.index, y=df_stats['Training Loss'],
                    mode='lines+markers',
                    name='Training Loss'))
    fig.add_trace(go.Scatter(x=df_stats.index, y=df_stats['Valid. Loss'],
                    mode='lines+markers',
                    name='Validation Loss'))
    fig.show()
"""
# <p style="color:#159364; font-family:cursive;">LEARNING RATE SCHEDULER </center></p>
"""
"""
**LINEAR SCHEDULE WITH WARMUP**
"""
# Defining LR Scheduler
scheduler = get_linear_schedule_with_warmup(
    optimizer, 
    num_warmup_steps=0, 
    num_training_steps=len(train_dataloader)*CONFIG.epochs
)
lrs = []
for epoch in range(1, CONFIG.epochs + 1):
    if scheduler is not None:
        scheduler.step()
    lrs.append(optimizer.param_groups[0]["lr"])
layout = go.Layout(template= "plotly_dark",title='Learning_rate')
fig = go.Figure(layout=layout)

fig.add_trace(go.Scatter(x=list(range(10)), y=lrs,
                    mode='lines+markers',
                    name='Learning_rate'))
fig.show()
"""
**RUN ON THE SET SCHEDULER**
"""
df1=run(optimizer,scheduler)
Visualizations(df1)
"""

# <p><center style="color:#159364; font-family:cursive;">ANALYZE LEARNING RATE VARIATION IN OTHER LEARNING RATE SCHEDULERS</center></p>
"""
"""
# 1.Reduce LR On Plateau:
"""
"""
* Reduces learning rate when a metric has stopped improving. Models often benefit from reducing the learning rate by a factor of 2-10 once learning stagnates. This scheduler reads a metrics quantity and if no improvement is seen for a ‘patience’ number of epochs, the learning rate is reduced.
* In this scheduler,the scheduler.step() requires the loss argument to identify any stagnation in improvement of metric,so we changed the run function
"""
scheduler=torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.01, patience=1, threshold=0.0001, threshold_mode='rel', cooldown=0, min_lr=0, eps=1e-08, verbose=False)
def run_model(optimizer,scheduler,epochs):
    set_seed(40)
    scaler=CONFIG.scaler
    training_stats = []
    total_t0 = time.time()

    for epoch_i in range(0, epochs):
        print("")
        print('======== Epoch {:} / {:} ========'.format(epoch_i + 1, epochs))
        print('Training...')
        t0 = time.time()
        total_train_loss = 0
        data_size=0
        model.train()
        for step, batch in enumerate(train_dataloader):
            tr_loss=[]
            b_input_ids = batch['ids'].to(CONFIG.device)
            b_token_type_ids=batch['token_type_ids'].to(CONFIG.device)
            b_input_mask = batch['mask'].to(CONFIG.device)
            b_labels = batch['target'].to(CONFIG.device)
            batch_size = b_input_ids.size(0)
            model.zero_grad() 
            with amp.autocast(enabled=True):
                output= model(b_input_ids,attention_mask=b_input_mask)          
                output=output["logits"].squeeze(-1)
                loss = loss_fn(output,b_labels)
                tr_loss.append(loss.item()/len(output))
            scheduler.step(loss)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        avg_train_loss = np.mean(tr_loss)    
        training_time = format_time(time.time() - t0)

        gc.collect()
        print("")
        print("  Average training loss: {0:.2f}".format(avg_train_loss))
        print("  Training epoch took: {:}".format(training_time))
        print("")
        print("Running Validation...")

        t0 = time.time()
        model.eval()
        val_loss = 0
        allpreds = []
        alltargets = []

        for batch in validation_dataloader:
            losses = []
            with torch.no_grad():
                device=CONFIG.device
                ids = batch["ids"].to(device)
                mask = batch["mask"].to(device)
                tokentype = batch["token_type_ids"].to(device)
                output = model(ids,mask)
                output = output["logits"].squeeze(-1)
                target = batch["target"].to(device)
                loss = loss_fn(output,target)
                losses.append(loss.item()/len(output))
                allpreds.append(output.detach().cpu().numpy())
                alltargets.append(target.detach().squeeze(-1).cpu().numpy())
        allpreds = np.concatenate(allpreds)
        alltargets = np.concatenate(alltargets)
        losses = np.mean(losses)
        gc.collect() 
        validation_time = format_time(time.time() - t0)
        print("  Validation Loss: {0:.2f}".format(losses))
        print("  Validation took: {:}".format(validation_time))
        training_stats.append(
        {
            'epoch': epoch_i + 1,
            'Training Loss': avg_train_loss,
            'lr':optimizer.param_groups[0]['lr'],
            'Valid. Loss': losses,
            'Training Time': training_time,
            'Validation Time': validation_time
        }
    ) 
    print("")
    print("Training complete!")
    return training_stats
    


    
"""
# LAMBDA LR
"""
"""
Sets the learning rate of each parameter group to the initial lr times a given function(here it is 0.65^epoch). When last_epoch=-1, sets initial lr as lr.
"""
optimizer = AdamW(optimizer_parameters, lr=CONFIG.learning_rate)
lambda1 = lambda epoch: 0.65 ** epoch
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda1)
lrs = []
for i in range(10):
    optimizer.step()
    lrs.append(optimizer.param_groups[0]["lr"])
    scheduler.step()
layout = go.Layout(template= "plotly_dark",title='Learning_rate')
fig = go.Figure(layout=layout)

fig.add_trace(go.Scatter(x=list(range(10)), y=lrs,
                    mode='lines+markers',
                    name='Learning_rate'))
fig.show()
"""
# MULTIPLICATIVE LR
"""
"""
Multiplies the learning rate of each parameter group by the factor given in the specified function. When last_epoch=-1, sets initial lr as lr.
"""
optimizer = AdamW(optimizer_parameters, lr=CONFIG.learning_rate)

lmbda = lambda epoch: 0.65 ** epoch
scheduler = torch.optim.lr_scheduler.MultiplicativeLR(optimizer, lr_lambda=lmbda)
lrs = []

for i in range(10):
    optimizer.step()
    lrs.append(optimizer.param_groups[0]["lr"])
    scheduler.step()

layout = go.Layout(template= "plotly_dark",title='Learning_rate')
fig = go.Figure(layout=layout)

fig.add_trace(go.Scatter(x=list(range(10)), y=lrs,
                    mode='lines+markers',
                    name='Learning_rate'))
fig.show()

"""
# STEP LR
"""
"""
Decays the learning rate of each parameter group by gamma every step_size epochs
"""

scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.1)
lrs = []

for i in range(10):
    optimizer.step()
    lrs.append(optimizer.param_groups[0]["lr"])
    scheduler.step()

layout = go.Layout(template= "plotly_dark",title='Learning_rate')
fig = go.Figure(layout=layout)

fig.add_trace(go.Scatter(x=list(range(10)), y=lrs,
                    mode='lines+markers',
                    name='Learning_rate'))
fig.show()

"""
# Multi Step LR
"""
"""
Decays the learning rate of each parameter group by gamma once the number of epoch reaches one of the milestones.
"""
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[6,8,9], gamma=0.1)
lrs = []

for i in range(10):
    optimizer.step()
    lrs.append(optimizer.param_groups[0]["lr"])
    scheduler.step()

layout = go.Layout(template= "plotly_dark",title='Learning_rate')
fig = go.Figure(layout=layout)

fig.add_trace(go.Scatter(x=list(range(10)), y=lrs,
                    mode='lines+markers',
                    name='Learning_rate'))
fig.show()

"""
# ExponentialLR
"""
"""
Decays the learning rate of each parameter group by gamma every epoch.
"""
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.1)
lrs = []
for i in range(10):
    optimizer.step()
    lrs.append(optimizer.param_groups[0]["lr"])
    scheduler.step()

layout = go.Layout(template= "plotly_dark",title='Learning_rate')
fig = go.Figure(layout=layout)

fig.add_trace(go.Scatter(x=list(range(10)), y=lrs,
                    mode='lines+markers',
                    name='Learning_rate'))
fig.show()

"""
# Cosine Annealing LR
"""
"""
Sets the learning rate of each parameter group using a cosine annealing schedule.
If the learning rate is set solely by this scheduler, the learning rate at each step becomes:

$$
\eta_{t}=\eta_{\min }+\frac{1}{2}\left(\eta_{\max }-\eta_{\min }\right)\left(1+\cos \left(\frac{T_{c u r}}{T_{\max }} \pi\right)\right)
$$

It has been proposed in SGDR: Stochastic Gradient Descent with Warm Restarts. Note that this only implements the cosine annealing part of SGDR, and not the restarts.https://arxiv.org/abs/1608.03983
"""
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10, eta_min=0)
lrs = []
for i in range(100):
    optimizer.step()
    lrs.append(optimizer.param_groups[0]["lr"])
    scheduler.step()

layout = go.Layout(template= "plotly_dark",title='Learning_rate')
fig = go.Figure(layout=layout)

fig.add_trace(go.Scatter(x=list(range(100)), y=lrs,
                    mode='lines+markers',
                    name='Learning_rate'))
fig.show()

"""
# CYCLIC LR
Sets the learning rate of each parameter group according to cyclical learning rate policy (CLR). The policy cycles the learning rate between two boundaries with a constant frequency, as detailed in the paper [Cyclical Learning Rates for Training Neural Networks](https://arxiv.org/abs/1506.01186). The distance between the two boundaries can be scaled on a per-iteration or per-cycle basis.

Cyclical learning rate policy changes the learning rate after every batch. step should be called after a batch has been used for training.
This class has three built-in policies, as put forth in the paper:

* **“triangular”**: A basic triangular cycle without amplitude scaling.
* **“triangular2”**: A basic triangular cycle that scales initial amplitude by half each cycle.
* **“exp_range”**: A cycle that scales initial amplitude by \text{gamma}^{\text{cycle iterations}}gamma cycle iterations at each cycle iteration.
"""
"""
# CYCLIC LR:TRIANGULAR
"""
optimizer = torch.optim.SGD(model.parameters(), lr=100)
scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, base_lr=0.001, max_lr=0.1,step_size_up=5,mode="triangular")
lrs = []


for i in range(100):
    optimizer.step()
    lrs.append(optimizer.param_groups[0]["lr"])
    scheduler.step()

layout = go.Layout(template= "plotly_dark",title='Learning_rate')
fig = go.Figure(layout=layout)

fig.add_trace(go.Scatter(x=list(range(100)), y=lrs,
                    mode='lines+markers',
                    name='Learning_rate'))
fig.show()

"""
# CYCLIC LR:TRIANGUALR2
"""
optimizer = torch.optim.SGD(model.parameters(), lr=100)
scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, base_lr=0.001, max_lr=0.1,step_size_up=5,mode="triangular2")
lrs = []


for i in range(100):
    optimizer.step()
    lrs.append(optimizer.param_groups[0]["lr"])
    scheduler.step()

layout = go.Layout(template= "plotly_dark",title='Learning_rate')
fig = go.Figure(layout=layout)

fig.add_trace(go.Scatter(x=list(range(100)), y=lrs,
                    mode='lines+markers',
                    name='Learning_rate'))
fig.show()

"""
**RUNNING OUR MODEL WITH THIS SCHEDULER**
"""
df9=run(optimizer,scheduler)

Visualizations(df9)
"""
# CYCLIC LR:EXPONENTIAL RANGE
"""
optimizer = torch.optim.SGD(model.parameters(), lr=100)
scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, base_lr=0.001, max_lr=0.1,step_size_up=5,mode="exp_range",gamma=0.85)
lrs = []


for i in range(100):
    optimizer.step()
    lrs.append(optimizer.param_groups[0]["lr"])
    scheduler.step()

layout = go.Layout(template= "plotly_dark",title='Learning_rate')
fig = go.Figure(layout=layout)

fig.add_trace(go.Scatter(x=list(range(100)), y=lrs,
                    mode='lines+markers',
                    name='Learning_rate'))
fig.show()

"""
# ONE CYCLE LR:COSINE ANNEALING (DEFAULT)
"""
"""
* Sets the learning rate of each parameter group according to the 1 cycle learning rate policy.
* The 1 cycle policy anneals the learning rate from an initial learning rate to some maximum learning rate and then from that maximum learning rate to some minimum learning rate much lower than the initial learning rate.
"""
optimizer = torch.optim.SGD(model.parameters(), lr=100)
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.1, steps_per_epoch=10, epochs=10)
lrs = []


for i in range(100):
    optimizer.step()
    lrs.append(optimizer.param_groups[0]["lr"])
    scheduler.step()

layout = go.Layout(template= "plotly_dark",title='Learning_rate')
fig = go.Figure(layout=layout)

fig.add_trace(go.Scatter(x=list(range(100)), y=lrs,
                    mode='lines+markers',
                    name='Learning_rate'))
fig.show()

"""
# ONE CYCLE LR:LINEAR ANNEALING
"""
optimizer = torch.optim.SGD(model.parameters(), lr=100)
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.1, steps_per_epoch=10, epochs=10,anneal_strategy='linear')
lrs = []


for i in range(100):
    optimizer.step()
    lrs.append(optimizer.param_groups[0]["lr"])
    scheduler.step()
layout = go.Layout(template= "plotly_dark",title='Learning_rate')
fig = go.Figure(layout=layout)

fig.add_trace(go.Scatter(x=list(range(100)), y=lrs,
                    mode='lines+markers',
                    name='Learning_rate'))
fig.show()

"""
# COSINE ANNEALING WARM RESTARTS
"""
"""
Set the learning rate of each parameter group using a cosine annealing schedule, and restarts after Ti epochs.


$$
\eta_{t}=\eta_{\min }+\frac{1}{2}\left(\eta_{\max }-\eta_{\min }\right)\left(1+\cos \left(\frac{T_{\operatorname{cur}}}{T_{i}} \pi\right)\right)
$$

"""
lr_sched = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=1, eta_min=0.001, last_epoch=-1)


lrs = []

for i in range(100):
    lr_sched.step()
    lrs.append(
        optimizer.param_groups[0]["lr"]
    )

layout = go.Layout(template= "plotly_dark",title='Learning_rate')
fig = go.Figure(layout=layout)

fig.add_trace(go.Scatter(x=list(range(100)), y=lrs,
                    mode='lines+markers',
                    name='Learning_rate'))
fig.show()

lr_sched = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2, eta_min=0.01, last_epoch=-1)


lrs = []

for i in range(300):
    lr_sched.step()
    lrs.append(
        optimizer.param_groups[0]["lr"]
    )

layout = go.Layout(template= "plotly_dark",title='Learning_rate')
fig = go.Figure(layout=layout)

fig.add_trace(go.Scatter(x=list(range(300)), y=lrs,
                    mode='lines+markers',
                    name='Learning_rate'))
fig.show()

"""
![Upvote!](https://img.shields.io/badge/Upvote-If%20you%20like%20my%20work-07b3c8?style=for-the-badge&logo=kaggle)

"""