Checkpoint with optimizer state in PyTorch - Practice Problems & Coding Challenges

Challenge - 5 Problems

❓ Predict Output

intermediate

What is the output of this PyTorch checkpoint loading code?

Consider the following PyTorch code that saves and loads a model checkpoint including the optimizer state. What will be printed after loading?

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim

model = nn.Linear(2, 1)
optimizer = optim.SGD(model.parameters(), lr=0.1)

# Simulate one optimizer step
optimizer.zero_grad()
output = model(torch.tensor([[1.0, 2.0]]))
loss = output.sum()
loss.backward()
optimizer.step()

# Save checkpoint
checkpoint = {'model_state': model.state_dict(), 'optimizer_state': optimizer.state_dict()}
torch.save(checkpoint, 'checkpoint.pth')

# Create new model and optimizer
model2 = nn.Linear(2, 1)
optimizer2 = optim.SGD(model2.parameters(), lr=0.1)

# Load checkpoint
loaded = torch.load('checkpoint.pth')
model2.load_state_dict(loaded['model_state'])
optimizer2.load_state_dict(loaded['optimizer_state'])

# Check optimizer state keys
print(sorted(optimizer2.state_dict().keys()))

A ['param_groups', 'state']

B ['model_state', 'optimizer_state']

C ['weights', 'biases']

D ['learning_rate', 'momentum']

💡 Hint

Look at what keys are stored inside optimizer.state_dict() in PyTorch.

The optimizer.state_dict() returns a dictionary with two keys: 'state' which holds the internal state for each parameter (like momentum buffers), and 'param_groups' which holds parameter group info including learning rates.

❓ Model Choice

intermediate

Which optimizer state is necessary to save for resuming training exactly?

When saving a checkpoint to resume training later without losing optimizer progress, which part of the optimizer must be saved?

A Only the model's parameters

B Only the gradients of the model parameters

C Only the optimizer's hyperparameters like learning rate and momentum

D The optimizer's internal state (momentum buffers, etc.) and parameter groups

💡 Hint

Think about what the optimizer uses internally to update parameters beyond just hyperparameters.

To resume training exactly, you must save the optimizer's internal state (like momentum buffers) and parameter groups. This allows the optimizer to continue updating parameters as if training was never interrupted.

❓ Hyperparameter

advanced

What happens if you load an optimizer state with a different learning rate than the current optimizer?

Suppose you saved an optimizer state with learning rate 0.01 but now you create a new optimizer with learning rate 0.001 and load the saved state. What learning rate will the optimizer use after loading?

A It will raise an error due to mismatch

B It will keep 0.001 from the new optimizer ignoring the loaded state

C It will use 0.01 from the loaded state, overriding the new optimizer's setting

D It will average 0.01 and 0.001 and use 0.0055

💡 Hint

Loading optimizer state_dict overwrites all parameter groups including learning rates.

When loading optimizer state_dict, the parameter groups including learning rates are restored from the saved state, overriding any new settings.

🔧 Debug

advanced

Why does this checkpoint loading code cause a runtime error?

Given this code snippet, why does loading the optimizer state cause a runtime error? ```python model = nn.Linear(3, 2) optimizer = optim.Adam(model.parameters(), lr=0.01) checkpoint = torch.load('checkpoint.pth') model.load_state_dict(checkpoint['model_state']) optimizer.load_state_dict(checkpoint['optimizer_state']) ``` Assume the checkpoint was saved from a model with 2 input features instead of 3.

A The model state dict keys do not match, causing optimizer load to fail

B The optimizer state keys do not match because model parameters changed shape

C The checkpoint file is corrupted

D The optimizer type must be the same but it is different

💡 Hint

Check if model parameter shapes match between saved and current model.

If the model parameters have different shapes, the optimizer state dict cannot map saved states to current parameters, causing a runtime error when loading optimizer state.

🧠 Conceptual

expert

Why is saving optimizer state important for training with adaptive optimizers?

Adaptive optimizers like Adam keep internal statistics (e.g., running averages of gradients). Why is saving and restoring the optimizer state critical when resuming training with such optimizers?

A Because internal statistics affect parameter updates and losing them changes training dynamics

B Because optimizer state controls the learning rate scheduler

C Because optimizer state contains the training data used so far

D Because without optimizer state, the model weights cannot be restored

💡 Hint

Think about what adaptive optimizers use internally to adjust updates.

Adaptive optimizers rely on internal running averages of gradients and squared gradients to adjust updates. Losing this state resets these statistics, changing training behavior and possibly harming convergence.

Practice

1. What is the main reason to save the optimizer state along with the model in a PyTorch checkpoint?

easy

A. To speed up the model's inference time

B. To reduce the model size on disk

C. To resume training with the same learning rate and momentum settings

D. To convert the model to a different format

Solution

1. Step 1: Understand what optimizer state contains

   The optimizer state includes parameters like learning rate, momentum, and other variables that affect training progress.

2. Step 2: Reason why saving optimizer state is important

   Saving the optimizer state allows training to resume exactly where it left off, preserving these settings.

3. Final Answer:

   To resume training with the same learning rate and momentum settings -> Option C

4. Quick Check:

   Optimizer state saves training settings = C [OK]

Hint: Optimizer state saves training progress settings [OK]

Common Mistakes:

• Thinking optimizer state reduces model size
• Confusing optimizer state with model weights
• Believing optimizer state affects inference speed

2. Which of the following is the correct way to save a checkpoint including model and optimizer states in PyTorch?

easy

A. torch.save(model, 'checkpoint.pth')

B. torch.save({'model': model.state_dict(), 'optimizer': optimizer.state_dict()}, 'checkpoint.pth')

C. torch.save(optimizer, 'checkpoint.pth')

D. torch.save({'model': model, 'optimizer': optimizer}, 'checkpoint.pth')

Solution

1. Step 1: Identify correct saving method for states

   PyTorch recommends saving state_dict() of model and optimizer for checkpoints.

2. Step 2: Check each option

   torch.save({'model': model.state_dict(), 'optimizer': optimizer.state_dict()}, 'checkpoint.pth') saves state_dict() of both model and optimizer in a dictionary, which is correct.

3. Final Answer:

   torch.save({'model': model.state_dict(), 'optimizer': optimizer.state_dict()}, 'checkpoint.pth') -> Option B

4. Quick Check:

   Save state_dict() for model and optimizer = B [OK]

Hint: Save state_dict() of model and optimizer in dict [OK]

Common Mistakes:

• Saving full model object instead of state_dict
• Saving optimizer object directly
• Not saving optimizer state at all

3. Given this code snippet, what will be printed?

import torch
import torch.nn as nn
import torch.optim as optim

model = nn.Linear(2, 1)
optimizer = optim.SGD(model.parameters(), lr=0.1)

# Save checkpoint
checkpoint = {'model': model.state_dict(), 'optimizer': optimizer.state_dict()}
torch.save(checkpoint, 'cp.pth')

# Load checkpoint
loaded = torch.load('cp.pth')
optimizer.load_state_dict(loaded['optimizer'])
print(optimizer.param_groups[0]['lr'])

medium

A. 0.1

B. 0.01

C. 1.0

D. Error: optimizer state not loaded

Solution

1. Step 1: Understand optimizer initialization

   Optimizer is created with learning rate 0.1 and saved in checkpoint.

2. Step 2: Loading optimizer state restores learning rate

   Loading optimizer state_dict sets learning rate back to 0.1.

3. Final Answer:

   0.1 -> Option A

4. Quick Check:

   Loaded optimizer lr = 0.1 [OK]

Hint: Loaded optimizer keeps saved learning rate [OK]

Common Mistakes:

• Assuming learning rate resets to default
• Forgetting to load optimizer state
• Confusing model and optimizer states

4. You saved a checkpoint with model and optimizer states but when loading, training behaves as if optimizer settings are lost. What is the most likely mistake?

medium

A. Not calling optimizer.load_state_dict() after loading checkpoint

B. Saving model.state_dict() instead of model

C. Using torch.save() instead of torch.load()

D. Not setting model.eval() before saving

Solution

1. Step 1: Identify cause of lost optimizer settings

   If optimizer state is not loaded, training uses default optimizer settings.

2. Step 2: Check common mistakes

   Not calling optimizer.load_state_dict() after loading checkpoint causes this issue.

3. Final Answer:

   Not calling optimizer.load_state_dict() after loading checkpoint -> Option A

4. Quick Check:

   Load optimizer state to keep settings = D [OK]

Hint: Always load optimizer state after loading checkpoint [OK]

Common Mistakes:

• Saving full model instead of state_dict
• Confusing torch.save and torch.load usage
• Setting model.eval() affects inference, not optimizer

5. You want to save a checkpoint that allows resuming training exactly, including epoch number and best loss so far. Which is the best way to structure the checkpoint dictionary?

hard

A. {'epoch': epoch, 'model': model.state_dict()}

B. {'model': model, 'optimizer': optimizer, 'epoch': epoch}

C. {'model_state': model.state_dict(), 'loss': best_loss}

D. {'epoch': epoch, 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'best_loss': best_loss}

Solution

1. Step 1: Identify required checkpoint components

   To resume training exactly, save epoch, model state, optimizer state, and best loss.

2. Step 2: Evaluate options

   {'epoch': epoch, 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'best_loss': best_loss} includes all required keys with correct state_dict() usage.

3. Final Answer:

   {'epoch': epoch, 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'best_loss': best_loss} -> Option D

4. Quick Check:

   Save epoch, model, optimizer, loss in checkpoint = A [OK]

Hint: Include epoch, model, optimizer, and loss in checkpoint dict [OK]

Common Mistakes:

• Saving full model or optimizer objects
• Omitting optimizer state
• Not saving epoch or loss for training resume