Checkpoint with optimizer state in PyTorch

Introduction

Saving a checkpoint with the optimizer state lets you pause and continue training later without losing progress.

You want to stop training and resume later without starting over.

You want to save your model and optimizer to recover from crashes.

You want to try different training settings starting from the same point.

You want to share your trained model and optimizer state with others.

Syntax

PyTorch

torch.save({
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'epoch': epoch,
    'loss': loss_value
}, PATH)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss_value = checkpoint['loss']

Use model.state_dict() and optimizer.state_dict() to get their states.

Loading optimizer state restores learning rates and momentum for smooth training continuation.

Examples

Saves model and optimizer states to a file named checkpoint.pth.

PyTorch

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

Loads the saved states back into model and optimizer.

PyTorch

checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

Also saves the current epoch number to resume training from the right place.

PyTorch

torch.save({'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'epoch': current_epoch}, 'checkpoint.pth')

Sample Model

This code trains a simple model for one step, saves the model and optimizer states along with epoch and loss, then loads them back and prints the saved epoch and loss.

PyTorch

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

# Simple model
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 1)
    def forward(self, x):
        return self.linear(x)

# Create model and optimizer
model = SimpleNet()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# Dummy data
inputs = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
targets = torch.tensor([[1.0], [2.0]])

# Loss function
criterion = nn.MSELoss()

# Training step
model.train()
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()

# Save checkpoint
checkpoint_path = 'checkpoint.pth'
torch.save({
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'epoch': 1,
    'loss': loss.item()
}, checkpoint_path)

# Load checkpoint
checkpoint = torch.load(checkpoint_path)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
loaded_epoch = checkpoint['epoch']
loaded_loss = checkpoint['loss']

print(f"Loaded epoch: {loaded_epoch}")
print(f"Loaded loss: {loaded_loss:.4f}")

Important Notes

Always save both model and optimizer states to continue training smoothly.

Include extra info like epoch and loss to track training progress.

Use torch.load(PATH, map_location=torch.device('cpu')) if loading on CPU from GPU-trained model.

Summary

Checkpoint saves model and optimizer states to pause and resume training.

Loading optimizer state restores training settings like learning rate.

Include epoch and loss in checkpoint to track training progress.

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