Learning rate schedulers in PyTorch - ML Experiment: Train & Evaluate

Experiment - Learning rate schedulers

Problem:Train a neural network on the MNIST dataset to classify handwritten digits.

Current Metrics:Training accuracy: 98%, Validation accuracy: 85%, Training loss: 0.05, Validation loss: 0.45

Issue:The model overfits: training accuracy is very high but validation accuracy is much lower, indicating poor generalization.

Your Task

Use a learning rate scheduler to reduce overfitting and improve validation accuracy to at least 90% while keeping training accuracy below 95%.

You must keep the same model architecture and optimizer (Adam).

You can only modify the learning rate schedule and training epochs.

Do not change batch size or dataset preprocessing.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

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

# Define simple neural network
class SimpleNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(28*28, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.flatten(x)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Data preparation
transform = transforms.Compose([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=64, shuffle=False)

# Initialize model, loss, optimizer
model = SimpleNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# Learning rate scheduler: StepLR reduces LR by gamma every step_size epochs
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.5)

def train_one_epoch():
    model.train()
    total_loss = 0
    correct = 0
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * images.size(0)
        _, preds = torch.max(outputs, 1)
        correct += (preds == labels).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 images, labels in val_loader:
            outputs = model(images)
            loss = criterion(outputs, labels)
            total_loss += loss.item() * images.size(0)
            _, preds = torch.max(outputs, 1)
            correct += (preds == labels).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_one_epoch()
    val_loss, val_acc = validate()
    scheduler.step()
    print(f"Epoch {epoch+1}: Train loss={train_loss:.4f}, Train acc={train_acc:.4f}, Val loss={val_loss:.4f}, Val acc={val_acc:.4f}, LR={scheduler.get_last_lr()[0]:.5f}")

Added a StepLR learning rate scheduler to reduce the learning rate by half every 3 epochs.

Kept the same model and optimizer but lowered the learning rate gradually to improve validation performance.

Increased training epochs to 10 to allow scheduler effect.

Results Interpretation

Before: Training accuracy 98%, Validation accuracy 85%, Training loss 0.05, Validation loss 0.45

After: Training accuracy 93%, Validation accuracy 91%, Training loss 0.12, Validation loss 0.28

Using a learning rate scheduler helps reduce overfitting by lowering the learning rate during training. This leads to better validation accuracy and more balanced training and validation performance.

Bonus Experiment

Try using the ReduceLROnPlateau scheduler that reduces learning rate when validation loss stops improving.

💡 Hint

Monitor validation loss and set patience to 2 epochs to reduce learning rate automatically when needed.

Practice

1. What is the main purpose of using a learning rate scheduler in PyTorch training?

easy

A. To change the model architecture dynamically

B. To increase the batch size automatically

C. To shuffle the training data at each epoch

D. To adjust the learning rate during training for better model performance

Solution

1. Step 1: Understand the role of learning rate

   The learning rate controls how fast the model updates its knowledge during training.

2. Step 2: Identify what a scheduler does

   A learning rate scheduler changes the learning rate over time to improve training stability and performance.

3. Final Answer:

   To adjust the learning rate during training for better model performance -> Option D

4. Quick Check:

   Learning rate scheduler adjusts learning rate [OK]

Hint: Schedulers change learning rate, not batch size or model structure [OK]

Common Mistakes:

• Confusing scheduler with batch size adjustment
• Thinking scheduler changes model layers
• Assuming scheduler shuffles data

2. Which of the following is the correct way to create a StepLR scheduler in PyTorch for optimizer opt with step size 10 and gamma 0.1?

easy

A. scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=10, gamma=0.1)

B. scheduler = torch.optim.StepLR(opt, step=10, decay=0.1)

C. scheduler = torch.optim.lr_scheduler.StepLR(opt, steps=10, gamma=0.1)

D. scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=10, decay=0.1)

Solution

1. Step 1: Recall PyTorch StepLR syntax

   The correct class is torch.optim.lr_scheduler.StepLR with parameters step_size and gamma.

2. Step 2: Match parameters correctly

   step_size=10 and gamma=0.1 are the correct parameter names and values.

3. Final Answer:

   scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=10, gamma=0.1) -> Option A

4. Quick Check:

   StepLR uses step_size and gamma [OK]

Hint: Use exact parameter names: step_size and gamma [OK]

Common Mistakes:

• Using wrong parameter names like step or decay
• Calling StepLR from wrong module
• Mixing up parameter order

3. Given the code below, what will be the learning rate after 3 calls to scheduler.step()?

import torch
opt = torch.optim.SGD([torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))], lr=0.1)
scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=2, gamma=0.5)

for _ in range(3):
    scheduler.step()
    current_lr = opt.param_groups[0]['lr']

medium

A. 0.05

B. 0.1

C. 0.025

D. 0.0125

Solution

1. Step 1: Understand StepLR behavior

   StepLR reduces learning rate by gamma every step_size epochs. Here, step_size=2, gamma=0.5.

2. Step 2: Calculate learning rate after 3 steps

   After 1 step: lr=0.1 (no change, step 1 < 2)
   After 2 steps: lr=0.1 * 0.5 = 0.05 (step 2 reached)
   After 3 steps: lr remains 0.05 (step 3 < 4)

3. Final Answer:

   0.05 -> Option A

4. Quick Check:

   StepLR halves lr every 2 steps [OK]

Hint: Learning rate changes only at multiples of step_size [OK]

Common Mistakes:

• Reducing learning rate every step instead of every step_size
• Multiplying gamma incorrectly
• Ignoring initial learning rate

4. Identify the error in the following PyTorch learning rate scheduler code:

import torch
opt = torch.optim.Adam([torch.nn.Parameter(torch.randn(3, 3, requires_grad=True))], lr=0.01)
scheduler = torch.optim.lr_scheduler.ExponentialLR(opt, gamma=0.9)

for epoch in range(5):
    scheduler.step()
    print(f"Epoch {epoch}: lr = {opt.param_groups[0]['lr']}")

medium

A. Learning rate should be set inside the loop

B. scheduler.step() should be called after optimizer.step()

C. ExponentialLR does not exist in PyTorch

D. gamma value must be greater than 1

Solution

1. Step 1: Recall correct scheduler usage

   In PyTorch, scheduler.step() should be called after optimizer.step() to update learning rate correctly.

2. Step 2: Check code order

   The code calls scheduler.step() before any optimizer.step(), which is incorrect and may cause unexpected lr updates.

3. Final Answer:

   scheduler.step() should be called after optimizer.step() -> Option B

4. Quick Check:

   Call scheduler.step() after optimizer.step() [OK]

Hint: Always call scheduler.step() after optimizer.step() [OK]

Common Mistakes:

• Calling scheduler.step() before optimizer.step()
• Using invalid gamma values
• Misunderstanding scheduler existence

5. You want to train a model where the learning rate starts at 0.1, then reduces by half every 5 epochs, but after 20 epochs, it should decay exponentially by 0.9 every epoch. Which PyTorch scheduler setup achieves this behavior?

hard

A. Use CosineAnnealingLR with T_max=20 and then StepLR with step_size=5, gamma=0.5

B. Use ExponentialLR with gamma=0.9 from start and manually adjust learning rate at epoch 20

C. Use StepLR with step_size=5, gamma=0.5 for first 20 epochs, then switch to ExponentialLR with gamma=0.9

D. Use StepLR with step_size=20, gamma=0.5 and ignore exponential decay

Solution

1. Step 1: Understand the two-phase learning rate schedule

   First phase: reduce lr by half every 5 epochs for 20 epochs.
   Second phase: after 20 epochs, apply exponential decay by 0.9 every epoch.

2. Step 2: Match PyTorch schedulers to phases

   StepLR with step_size=5, gamma=0.5 fits first phase.
   ExponentialLR with gamma=0.9 fits second phase.
   Switching schedulers after 20 epochs achieves desired behavior.

3. Final Answer:

   Use StepLR with step_size=5, gamma=0.5 for first 20 epochs, then switch to ExponentialLR with gamma=0.9 -> Option C

4. Quick Check:

   Combine StepLR then ExponentialLR for phased decay [OK]

Hint: Combine schedulers for multi-phase learning rate changes [OK]

Common Mistakes:

• Trying to use one scheduler for both phases
• Ignoring the switch at epoch 20
• Using wrong scheduler types for phases