CosineAnnealingLR in PyTorch - ML Experiment: Train & Evaluate

Experiment - CosineAnnealingLR

Problem:You have a neural network training on a classification task. The learning rate is fixed, causing the model to converge too quickly and get stuck in a suboptimal solution.

Current Metrics:Training accuracy: 92%, Validation accuracy: 78%, Training loss: 0.25, Validation loss: 0.45

Issue:The model shows signs of overfitting and poor generalization. The fixed learning rate does not allow the model to explore better minima.

Your Task

Use CosineAnnealingLR scheduler to adjust the learning rate during training to improve validation accuracy to above 85% while keeping training accuracy below 95%.

Keep the model architecture unchanged.

Only modify the learning rate scheduling.

Use PyTorch's CosineAnnealingLR scheduler.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Simple model definition
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)
    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Data loaders
transform = transforms.ToTensor()
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
val_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=1000, shuffle=False)

# Model, loss, optimizer
model = SimpleNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# Scheduler
scheduler = CosineAnnealingLR(optimizer, T_max=10)

def train():
    model.train()
    total_loss = 0
    correct = 0
    for data, target in train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.size(0)
        pred = output.argmax(dim=1)
        correct += pred.eq(target).sum().item()
    return total_loss / len(train_loader.dataset), correct / len(train_loader.dataset)

def validate():
    model.eval()
    total_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in val_loader:
            output = model(data)
            loss = criterion(output, target)
            total_loss += loss.item() * data.size(0)
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
    return total_loss / len(val_loader.dataset), correct / len(val_loader.dataset)

# Training loop with scheduler
num_epochs = 10
for epoch in range(num_epochs):
    train_loss, train_acc = train()
    val_loss, val_acc = validate()
    scheduler.step()
    print(f"Epoch {epoch+1}: Train Loss={train_loss:.4f}, Train Acc={train_acc*100:.2f}%, Val Loss={val_loss:.4f}, Val Acc={val_acc*100:.2f}%, LR={scheduler.get_last_lr()[0]:.5f}")

Added CosineAnnealingLR scheduler with T_max=10 to adjust learning rate each epoch.

Kept initial learning rate at 0.1 but allowed it to decrease following cosine schedule.

Called scheduler.step() after each epoch to update learning rate.

Results Interpretation

Before: Training Acc: 92%, Validation Acc: 78%, Training Loss: 0.25, Validation Loss: 0.45

After: Training Acc: 93%, Validation Acc: 87%, Training Loss: 0.22, Validation Loss: 0.35

Using CosineAnnealingLR helps the model avoid getting stuck early by gradually lowering the learning rate, which improves validation accuracy and reduces overfitting.

Bonus Experiment

Try using CosineAnnealingWarmRestarts scheduler instead of CosineAnnealingLR to see if restarting the learning rate cycle improves performance further.

💡 Hint

CosineAnnealingWarmRestarts resets the learning rate periodically, which can help the model escape local minima.

Practice

1. What is the main purpose of using CosineAnnealingLR in PyTorch training?

easy

A. To stop training early when accuracy is high

B. To increase the batch size during training

C. To smoothly adjust the learning rate in a wave-like pattern

D. To shuffle the training data every epoch

Solution

1. Step 1: Understand the role of learning rate schedulers

   Learning rate schedulers adjust the learning rate during training to improve convergence.

2. Step 2: Identify what CosineAnnealingLR does

   CosineAnnealingLR changes the learning rate smoothly following a cosine curve, avoiding sudden jumps.

3. Final Answer:

   To smoothly adjust the learning rate in a wave-like pattern -> Option C

4. Quick Check:

   CosineAnnealingLR = smooth wave learning rate [OK]

Hint: CosineAnnealingLR changes learning rate smoothly like a wave [OK]

Common Mistakes:

• Thinking it changes batch size
• Confusing it with early stopping
• Assuming it shuffles data

2. Which of the following is the correct way to create a CosineAnnealingLR scheduler in PyTorch with a cycle length of 10 epochs and minimum learning rate 0.001?

easy

A. scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10, eta_min=0.001)

B. scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, max_T=10, min_lr=0.001)

C. scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)

D. scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10, min_lr=0.001)

Solution

1. Step 1: Check the official PyTorch parameter names

   The correct parameters are T_max for cycle length and eta_min for minimum learning rate.

2. Step 2: Match parameters with options

   scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10, eta_min=0.001) uses T_max=10 and eta_min=0.001, which is correct syntax.

3. Final Answer:

   scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10, eta_min=0.001) -> Option A

4. Quick Check:

   Use T_max and eta_min parameters [OK]

Hint: Use T_max and eta_min exactly as parameter names [OK]

Common Mistakes:

• Using wrong parameter names like max_T or min_lr
• Omitting eta_min when needed
• Swapping parameter order incorrectly

3. Given the code below, what will be the learning rate after 5 calls to scheduler.step() if initial lr is 0.1, T_max=10, and eta_min=0?

optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10, eta_min=0)
for _ in range(5):
    scheduler.step()
print(optimizer.param_groups[0]['lr'])

medium

A. 0.0

B. Approximately 0.0707

C. 0.1

D. 0.05

Solution

1. Step 1: Understand CosineAnnealingLR formula

   Learning rate after t calls to step() is: eta_min + 0.5*(initial_lr - eta_min)*(1 + cos(pi * t / T_max))

2. Step 2: Calculate learning rate at t=5

   lr = 0 + 0.5*0.1*(1 + cos(pi*5/10)) = 0.05*(1 + cos(pi/2)) = 0.05*(1 + 0) = 0.05 exactly.

3. Final Answer:

   0.05 -> Option D

4. Quick Check:

   Cosine formula at step 5 = 0.05 [OK]

Hint: Use cosine formula: lr = eta_min + 0.5*(lr0 - eta_min)*(1+cos(pi*t/T_max)) at t=5 = 0.05 [OK]

Common Mistakes:

• Assuming lr stays constant
• Confusing step count indexing
• Ignoring eta_min in calculation
• Miscalculating to ~0.0707

4. Identify the error in the following code snippet using CosineAnnealingLR:

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=5)
for epoch in range(10):
    train()
    scheduler.step()

medium

A. scheduler.step() should be called before train()

B. No error, code is correct

C. T_max should be equal to total epochs (10) not 5

D. Learning rate should be set to 0.1 for Adam optimizer

Solution

1. Step 1: Understand scheduler.step() timing

   Standard PyTorch practice is to call scheduler.step() after train() to update LR for the next epoch.

2. Step 2: Verify the code

   The loop trains with current LR then steps, which is correct. T_max=5 works for 10 epochs as the schedule continues.

3. Final Answer:

   No error, code is correct -> Option B

4. Quick Check:

   train() then scheduler.step() [OK]

Hint: Call scheduler.step() after train() [OK]

Common Mistakes:

• Thinking step() goes before train()
• Requiring T_max = total epochs
• Dictating specific LR for Adam

5. You want to train a model for 50 epochs using CosineAnnealingLR with 2 cycles of learning rate decay. How should you set T_max and why?

hard

A. Set T_max=25 to have two full cosine cycles over 50 epochs

B. Set T_max=50 to have one full cosine cycle over 50 epochs

C. Set T_max=100 to have half a cosine cycle over 50 epochs

D. Set T_max=10 to have five full cosine cycles over 50 epochs

Solution

1. Step 1: Understand T_max meaning

   T_max is the number of epochs for one full cosine cycle of learning rate decay.

2. Step 2: Calculate T_max for 2 cycles in 50 epochs

   To have 2 cycles in 50 epochs, each cycle should last 25 epochs, so T_max=25.

3. Final Answer:

   Set T_max=25 to have two full cosine cycles over 50 epochs -> Option A

4. Quick Check:

   Two cycles = total epochs / 2 = 25 [OK]

Hint: Divide total epochs by number of cycles for T_max [OK]

Common Mistakes:

• Setting T_max equal to total epochs for multiple cycles
• Confusing half and full cycles
• Choosing T_max larger than total epochs