Warmup + Cosine Decay Schedule

hardPROLearning Rate SchedulersPyTorch Advanced

Problem Statement

Implement a learning rate schedule with linear warmup followed by cosine decay using LambdaLR.

Background

Modern training often uses warmup (gradually increasing LR) followed by decay. This prevents early training instability. LambdaLR takes a function that returns a multiplier for the initial LR.

Your Task

The starter code creates a model and optimizer with warmup/total epoch constants. Create a LambdaLR scheduler that linearly warms up for 5 epochs, then follows cosine decay for the remaining epochs.

Output Format

Returns a dictionary with "lr_history" (20 values), "warmup_end_lr", and "peak_lr".

Example

Input

None

Output

{'lr_history': [0.02, 0.04, 0.06, 0.08, 0.1, 0.0995, 0.0979, 0.0952, 0.0916, 0.087, 0.0816, 0.0755, 0.069, 0.062, 0.0549, 0.0477, 0.0407, 0.0341, 0.028, 0.0226], 'warmup_end_lr': 0.1, 'peak_lr': 0.1}

Reasoning

The function warmup_cosine_test() initializes a learning rate schedule with a base learning rate of 0.1 and warmup epochs of 5.
During the warmup phase (epochs 0-4), the learning rate increases linearly: $LR = 0.1 \cdot \frac{epoch + 1}{5}$ , resulting in learning rates of 0.02, 0.04, 0.06, 0.08, and 0.1.
After the warmup phase, the learning rate enters a cosine decay phase (epochs 5-19): $LR = 0.1 \cdot 0.5 \cdot (1 + \cos(\pi \cdot \frac{epoch - 5}{15}))$ , producing the remaining learning rates in the lr_history list.
The function records the learning rate history, warmup end learning rate, and peak learning rate, returning them as a dictionary with the specified values.

Constraints

Linear warmup for 5 epochs
Cosine decay for remaining 15 epochs
Use LambdaLR

Learn this in

PyTorch FundamentalsFoundations Neural NetworksML Study Plan

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 primerPREVIEW

Background Knowledge

The problem involves implementing a learning rate schedule with linear warmup followed by cosine decay using LambdaLR in PyTorch. The learning rate schedule is a crucial component in training deep neural networks, as it controls how quickly the model learns from the data. A warmup period is often used to prevent early training instability by gradually increasing the learning rate. After the warmup period, the learning rate is decayed using a cosine decay schedule, which helps the model converge to a better solution.

The LambdaLR scheduler in PyTorch allows users to define a custom learning rate schedule using a lambda function. This function takes the current epoch as input and returns the corresponding learning rate. In this problem, the lambda function is defined to increase the learning rate linearly during the warmup period and then decay it using a cosine schedule. The cosine decay schedule is based on the cosine function, which oscillates between -1 and 1. The learning rate is scaled by the cosine function to create a decay schedule that helps the model converge.

The key concepts involved in this problem are learning rate schedules, warmup, cosine decay, and LambdaLR. Understanding these concepts is essential to implementing the solution. The learning rate schedule is critical in controlling the training process, and the warmup and cosine decay schedules are commonly used techniques to improve the stability and convergence of the model.

Algorithm/Approach

The general approach to solving this problem involves defining a custom learning rate schedule using a lambda function and then using the LambdaLR scheduler to apply this schedule to the optimizer. The lambda function should take into account the current epoch and the warmup and decay periods to calculate the corresponding learning rate. The LambdaLR scheduler will then use this lambda function to update the learning rate at each epoch.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results