StepLR Scheduler

easyFREELearning Rate SchedulersPyTorch Advanced

Problem Statement

Use StepLR to decay the learning rate by a factor every N epochs.

Background

StepLR(optimizer, step_size, gamma) multiplies LR by gamma every step_size epochs. It's the simplest scheduling strategy.

Your Task

The starter code creates a model and SGD optimizer. Create a StepLR scheduler that halves the learning rate every 3 epochs.

The epoch loop that records LR history is pre-filled.

Output Format

Returns a dictionary with "lr_history" (10 values), "initial_lr", and "final_lr".

Example

Input

None

Output

{'lr_history': [0.1, 0.1, 0.1, 0.05, 0.05, 0.05, 0.025, 0.025, 0.025, 0.0125], 'initial_lr': 0.1, 'final_lr': 0.0125}

Reasoning

The function takes no arguments. It starts from a learning rate $\eta_0 = 0.1$ , a step size of $3$ epochs and a decay factor $\gamma = 0.5$ and records the rate at the start of each of 10 epochs, $e = 0, \dots, 9$ .
The rate is constant within each block of 3 epochs and is multiplied by $\gamma$ each time a block ends: $\eta_e = 0.1 \cdot 0.5^{\lfloor e/3 \rfloor}$ .
Epochs 0–2 have exponent 0, so $\eta = 0.1$ . Epochs 3–5 have exponent 1, so $\eta = 0.05$ . Epochs 6–8 give $0.1 \cdot 0.25 = 0.025$ , and epoch 9 gives $0.1 \cdot 0.125 = 0.0125$ .
Each value is recorded before that epoch's decay is applied, which is why the first entry is the untouched $0.1$ . Epoch 9 opens the fourth block; the next halving would come at epoch 12, outside the window.
initial_lr is the first entry, 0.1, and final_lr is the last, 0.0125. Notice the history is a staircase, not a smooth curve: a step size of 2 would give five two-epoch plateaus.

Constraints

step_size=3, gamma=0.5
Record LR before each scheduler.step()
10 epochs total

Learn this in

3D Gaussian Splatting3D Computer Vision

More from PyTorch Advanced

Implement SGD with Momentumhard Implement Custom SGD Optimizerhard Weight Decay for Specific Layersmedium Freeze Backbone, Train Headmedium Per-Layer Learning Rateseasy

Concept primer

Background Knowledge

The learning rate is a crucial hyperparameter in training neural networks. It controls how quickly the model learns from the training data. A high learning rate can lead to fast convergence but may also cause the model to overshoot the optimal solution, while a low learning rate can result in slow convergence. To address this issue, learning rate schedulers are used to adjust the learning rate during training. One such scheduler is the StepLR scheduler, which decays the learning rate by a factor at every specified interval.

The StepLR scheduler takes three main parameters: the optimizer, step_size, and gamma. The optimizer is the optimization algorithm used to update the model's parameters, such as SGD (Stochastic Gradient Descent). The step_size determines the interval at which the learning rate is decayed, and gamma is the factor by which the learning rate is multiplied at each interval. For example, if step_size is 3 and gamma is 0.5, the learning rate will be halved every 3 epochs.

Understanding how the StepLR scheduler works is essential to solving this problem. The scheduler's step() method is called at each epoch to update the learning rate. The learning rate at each epoch can be retrieved from the optimizer object. By simulating the training process and recording the learning rate at each epoch, we can analyze how the StepLR scheduler affects the learning rate over time.

Algorithm/Approach

The general approach to solving this problem involves creating a neural network model, an optimizer, and a StepLR scheduler. We then simulate the training process by calling the step() method of the scheduler at each epoch. At each epoch, we record the learning rate and store it in a list. Finally, we return a dictionary containing the learning rate history, initial learning rate, and final learning rate.

Step-by-Step Strategy

To solve this problem, follow these steps:

Import the necessary PyTorch modules, including nn and optim.
Create a simple neural network model using nn.Linear.
Create an optimizer object, such as SGD, with the specified learning rate.
Create a StepLR scheduler object with the specified step_size and gamma.
Simulate the training process by iterating over the specified number of epochs.
At each epoch, call the step() method of the scheduler to update the learning rate.
Record the learning rate at each epoch and store it in a list.
Calculate the initial and final learning rates.
Return a dictionary containing the learning rate history, initial learning rate, and final learning rate.

Common Pitfalls

When implementing this solution, watch out for the following:

Forgetting to call the step() method of the scheduler at each epoch.
Not recording the learning rate at each epoch.
Not using the correct step_size and gamma values.
Not handling the learning rate history correctly.

Time & Space Complexity

The time complexity of this solution is O(n), where n is the number of epochs. The space complexity is also O(n), as we need to store the learning rate history for each epoch. The solution involves simple iterations and dictionary operations, making it efficient in terms of time and space complexity.

Here is a sample code to get you started:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR

def step_lr_test():
 # Create a simple neural network model
 model = nn.Linear(2, 1)
 
 # Create an optimizer object
 optimizer = optim.SGD(model.parameters(), lr=0.1)
 
 # Create a StepLR scheduler object
 scheduler = StepLR(optimizer, step_size=3, gamma=0.5)
 
 # Initialize the learning rate history list
 lr_history = []
 
 # Simulate the training process
 for epoch in range(10):
 # Call the step() method of the scheduler
 scheduler.step()
 
 # Record the learning rate at each epoch
 lr = optimizer.param_groups['lr']
 lr_history.append(round(lr, 4))
 
 # Calculate the initial and final learning rates
 initial_lr = lr_history
 final_lr = lr_history[-1]
 
 # Return a dictionary containing the learning rate history, initial learning rate, and final learning rate
 return {
 "lr_history": lr_history,
 "initial_lr": initial_lr,
 "final_lr": final_lr
 }

solution.py

Python 3.13

Test results