Warmup strategies in PyTorch - Practice Problems & Coding Challenges

Challenge - 5 Problems

🧠 Conceptual

intermediate

Why use a learning rate warmup in training?

Which of the following best explains the main reason for using a learning rate warmup at the start of training a neural network?

A To prevent the model from diverging by starting with a very high learning rate

B To reduce the total training time by skipping early epochs

C To gradually increase the learning rate to avoid large updates that can destabilize early training

D To immediately reach the maximum learning rate for faster convergence

💡 Hint

Think about how sudden large updates affect a new model's weights.

A warmup gradually increases the learning rate from a small value to the target value to avoid large, unstable updates early in training.

❓ Predict Output

intermediate

Output of learning rate scheduler with warmup

What will be the learning rate printed at epoch 3 in the following PyTorch code?

PyTorch

import torch
from torch.optim import SGD
from torch.optim.lr_scheduler import LambdaLR

optimizer = SGD([torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))], lr=0.1)
warmup_epochs = 5

def lr_lambda(epoch):
    if epoch < warmup_epochs:
        return (epoch + 1) / warmup_epochs
    else:
        return 1.0

scheduler = LambdaLR(optimizer, lr_lambda=lr_lambda)

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

A Epoch 3: lr = 0.06

B Epoch 3: lr = 0.04

C Epoch 3: lr = 0.1

D Epoch 3: lr = 0.08

💡 Hint

Calculate (epoch + 1) / warmup_epochs * base_lr for epoch 3.

At epoch 3, lr_lambda returns (3+1)/5 = 0.8, so lr = 0.8 * 0.1 = 0.08.

❓ Model Choice

advanced

Choosing warmup strategy for transformer training

You are training a transformer model on a large text dataset. Which warmup strategy is most suitable to stabilize training and improve final accuracy?

A Linear warmup followed by cosine decay

B Step warmup with abrupt jumps in learning rate

C Exponential warmup with sudden drop after warmup

D No warmup, start with a fixed learning rate

💡 Hint

Consider smooth transitions in learning rate for stable training.

Linear warmup gradually increases the learning rate, and cosine decay smoothly reduces it, which is effective for transformers.

❓ Hyperparameter

advanced

Determining warmup steps for a training schedule

If you have a total of 100 epochs and want to use a warmup phase that lasts 10% of training, how many warmup steps should you set?

A 5 steps

B 10 steps

C 20 steps

D 50 steps

💡 Hint

Calculate 10% of 100 epochs.

10% of 100 epochs is 10 epochs, so warmup steps should be 10.

❓ Metrics

expert

Effect of warmup on training loss curve

During training with and without learning rate warmup, which difference in the training loss curve is expected?

A With warmup, loss starts higher and decreases smoothly; without warmup, loss fluctuates or spikes early

B With warmup, loss decreases abruptly; without warmup, loss decreases smoothly

C With warmup, loss increases steadily; without warmup, loss decreases steadily

D With warmup, loss remains constant; without warmup, loss decreases steadily

💡 Hint

Think about stability of updates at the start of training.

Warmup prevents large early updates, so loss decreases smoothly. Without warmup, large updates can cause spikes or fluctuations.

Practice

1. What is the main purpose of using a warmup strategy in PyTorch training?

easy

A. To immediately set the learning rate to its maximum value

B. To gradually increase the learning rate at the start of training

C. To decrease the learning rate throughout the entire training

D. To freeze model weights during the first epochs

Solution

1. Step 1: Understand what warmup means

   Warmup means starting with a low learning rate and increasing it slowly.

2. Step 2: Identify the goal of warmup

   This helps the model learn smoothly and avoid sudden big updates that can harm training.

3. Final Answer:

   To gradually increase the learning rate at the start of training -> Option B

4. Quick Check:

   Warmup = gradual learning rate increase [OK]

Hint: Warmup means slowly raising learning rate early [OK]

Common Mistakes:

• Thinking warmup immediately sets max learning rate
• Confusing warmup with learning rate decay
• Assuming warmup freezes model weights

2. Which PyTorch class is commonly used to implement a warmup learning rate schedule with a custom function?

easy

A. torch.optim.lr_scheduler.StepLR

B. torch.optim.lr_scheduler.ReduceLROnPlateau

C. torch.optim.lr_scheduler.LambdaLR

D. torch.optim.lr_scheduler.ExponentialLR

Solution

1. Step 1: Recall PyTorch schedulers for warmup

   LambdaLR allows defining a custom function to adjust learning rate.

2. Step 2: Match scheduler to warmup use

   Warmup needs a custom function to increase learning rate gradually, which LambdaLR supports.

3. Final Answer:

   torch.optim.lr_scheduler.LambdaLR -> Option C

4. Quick Check:

   Custom function scheduler = LambdaLR [OK]

Hint: LambdaLR lets you define custom learning rate changes [OK]

Common Mistakes:

• Choosing StepLR which uses fixed step decay
• Picking ReduceLROnPlateau which reacts to metrics
• Selecting ExponentialLR which decays exponentially

3. Given the following PyTorch code snippet, what will be the learning rate at epoch 3?

import torch
optimizer = torch.optim.SGD([torch.nn.Parameter(torch.randn(2, 2))], lr=0.1)

warmup_epochs = 5
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda epoch: min((epoch + 1) / warmup_epochs, 1))

for epoch in range(5):
    scheduler.step()
    print(f"Epoch {epoch+1} LR: {optimizer.param_groups[0]['lr']}")

medium

A. 0.06

B. 0.03

C. 0.10

D. 0.50

Solution

1. Step 1: Understand the lambda function for LR

   The lambda function returns (epoch+1)/5 until it reaches 1, scaling the base LR 0.1.

2. Step 2: Calculate LR at epoch 3 (0-based index)

   Epoch 3 means epoch=2, so LR factor = (2+1)/5 = 3/5 = 0.6. LR = 0.1 * 0.6 = 0.06.

3. Final Answer:

   0.06 -> Option A

4. Quick Check:

   Epoch 3 LR = 0.1 * 3/5 = 0.06 [OK]

Hint: Multiply base LR by (epoch+1)/warmup_epochs [OK]

Common Mistakes:

• Using epoch number directly without +1
• Confusing epoch index with count
• Assuming LR is constant during warmup

4. Identify the error in this PyTorch warmup scheduler code:

import torch
optimizer = torch.optim.Adam([torch.nn.Parameter(torch.randn(2, 2))], lr=0.01)
warmup_epochs = 3
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda epoch: epoch / warmup_epochs)

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

medium

A. The optimizer should be SGD, not Adam

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

C. The learning rate is not scaled by base LR

D. The lambda function returns 0 at epoch 0 causing zero LR

Solution

1. Step 1: Analyze lambda function behavior at epoch 0

   At epoch 0, lambda returns 0/3 = 0, so LR is zero, which stops learning initially.

2. Step 2: Understand why zero LR is a problem

   Zero LR means no weight updates, which can slow or stop training progress early.

3. Final Answer:

   The lambda function returns 0 at epoch 0 causing zero LR -> Option D

4. Quick Check:

   Epoch 0 LR = 0 causes no learning [OK]

Hint: Check if lambda returns zero at first epoch [OK]

Common Mistakes:

• Ignoring zero LR at start
• Thinking optimizer type causes error
• Confusing scheduler step order

5. You want to implement a warmup strategy that linearly increases the learning rate from 0 to 0.1 over 4 epochs, then keeps it constant. Which lr_lambda function correctly achieves this in PyTorch's LambdaLR?

hard

A. lambda epoch: min((epoch + 1) / 4, 1)

B. lambda epoch: epoch / 4

C. lambda epoch: 1 if epoch >= 4 else 0.1 * epoch

D. lambda epoch: (epoch + 1) * 0.1

Solution

1. Step 1: Understand the warmup goal

   Learning rate should increase linearly from 0 to 1 (scale factor) over 4 epochs, then stay at 1.

2. Step 2: Check each lambda function

   lambda epoch: min((epoch + 1) / 4, 1) uses min((epoch+1)/4, 1), which linearly increases from 0.25 to 1 by epoch 4, then stays at 1.

3. Final Answer:

   lambda epoch: min((epoch + 1) / 4, 1) -> Option A

4. Quick Check:

   Linear increase capped at 1 = lambda epoch: min((epoch + 1) / 4, 1) [OK]

Hint: Use min((epoch+1)/warmup_epochs, 1) for linear warmup [OK]

Common Mistakes:

• Not adding +1 to epoch causing zero start
• Multiplying by 0.1 inside lambda instead of base LR
• Using step function instead of linear increase