Why checkpointing preserves progress in PyTorch - Experiment to Prove It

Experiment - Why checkpointing preserves progress

Problem:You are training a neural network on a dataset, but the training takes a long time. If the training process stops unexpectedly, you lose all progress and must start over.

Current Metrics:Training accuracy: 85%, Validation accuracy: 80%, Training loss: 0.5, Validation loss: 0.6

Issue:No checkpointing is used, so if training is interrupted, all progress is lost and training must restart from scratch.

Your Task

Implement checkpointing to save the model and optimizer state during training so that training can resume from the last saved point without losing progress.

Use PyTorch's native checkpointing methods.

Save checkpoints every 2 epochs.

Ensure that after restarting from a checkpoint, training continues correctly.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# Simple dataset
X = torch.randn(1000, 10)
y = (X.sum(dim=1) > 0).long()
dataset = TensorDataset(X, y)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

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

model = SimpleNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# Checkpoint path
checkpoint_path = 'checkpoint.pth'

# Function to save checkpoint
def save_checkpoint(epoch, model, optimizer):
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict()
    }, checkpoint_path)

# Function to load checkpoint
def load_checkpoint(model, optimizer):
    checkpoint = torch.load(checkpoint_path)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    return checkpoint['epoch']

# Try to load checkpoint
start_epoch = 0
try:
    start_epoch = load_checkpoint(model, optimizer) + 1
    print(f'Resuming training from epoch {start_epoch}')
except FileNotFoundError:
    print('No checkpoint found, starting fresh training')

# Training loop
num_epochs = 10
for epoch in range(start_epoch, num_epochs):
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    for inputs, labels in dataloader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * inputs.size(0)
        _, predicted = outputs.max(1)
        correct += predicted.eq(labels).sum().item()
        total += labels.size(0)

    avg_loss = total_loss / total
    accuracy = correct / total * 100
    print(f'Epoch {epoch}: Loss={avg_loss:.4f}, Accuracy={accuracy:.2f}%')

    # Save checkpoint every 2 epochs
    if (epoch + 1) % 2 == 0:
        save_checkpoint(epoch, model, optimizer)
        print(f'Checkpoint saved at epoch {epoch + 1}')

Added functions to save and load checkpoints using torch.save and torch.load.

Saved model state, optimizer state, and current epoch in checkpoint.

Modified training loop to load checkpoint if available and resume training.

Saved checkpoint every 2 epochs to preserve progress.

Fixed checkpoint saved epoch print statement to show correct epoch number.

Results Interpretation

Before checkpointing: If training stops, all progress is lost and training restarts from epoch 0.

After checkpointing: Training resumes from the last saved epoch, preserving progress and saving time.

Checkpointing saves the model and optimizer states during training. This allows training to continue from the last saved point after interruptions, preventing loss of progress and saving time.

Bonus Experiment

Try implementing checkpointing that also saves the best model based on validation accuracy.

💡 Hint

Track validation accuracy each epoch and save a separate checkpoint only when validation accuracy improves.

Practice

1. What is the main reason for using checkpointing during PyTorch model training?

easy

A. To save the model's current state so training can resume later without loss

B. To speed up the training by skipping some layers

C. To reduce the size of the training dataset

D. To automatically tune hyperparameters during training

Solution

1. Step 1: Understand checkpointing purpose

   Checkpointing saves the model's current state including weights and optimizer info.

2. Step 2: Connect checkpointing to training progress

   This allows training to stop and resume later without losing progress.

3. Final Answer:

   To save the model's current state so training can resume later without loss -> Option A

4. Quick Check:

   Checkpointing = Save progress [OK]

Hint: Checkpointing means saving progress to continue later [OK]

Common Mistakes:

• Thinking checkpointing speeds up training
• Confusing checkpointing with data reduction
• Assuming checkpointing tunes hyperparameters

2. Which of the following is the correct PyTorch code snippet to save a checkpoint?

easy

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

B. torch.save(model.state_dict(), 'checkpoint.pth')

C. torch.load('checkpoint.pth')

D. optimizer.save('checkpoint.pth')

Solution

1. Step 1: Identify saving function

   torch.save() is used to save objects like model weights to a file.

2. Step 2: Check correct usage for saving model state

   model.state_dict() returns model weights; saving it with torch.save() is correct.

3. Final Answer:

   torch.save(model.state_dict(), 'checkpoint.pth') -> Option B

4. Quick Check:

   Save model weights = torch.save(state_dict) [OK]

Hint: Use torch.save with model.state_dict() to save checkpoint [OK]

Common Mistakes:

• Using torch.load instead of torch.save to save
• Trying to save optimizer with wrong method
• Confusing load_state_dict with saving

3. Given this code snippet, what will be printed after loading the checkpoint?

model = MyModel()
optimizer = torch.optim.Adam(model.parameters())
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state'])
optimizer.load_state_dict(checkpoint['optimizer_state'])
epoch = checkpoint['epoch']
print(epoch)

medium

A. An error because checkpoint keys are missing

B. The total number of model parameters

C. The optimizer learning rate

D. The epoch number saved in the checkpoint

Solution

1. Step 1: Understand checkpoint contents

   The checkpoint dictionary contains keys 'model_state', 'optimizer_state', and 'epoch'.

2. Step 2: Identify printed value

   Variable 'epoch' is assigned checkpoint['epoch'], so print(epoch) outputs the saved epoch number.

3. Final Answer:

   The epoch number saved in the checkpoint -> Option D

4. Quick Check:

   Print epoch from checkpoint = epoch number [OK]

Hint: Print shows saved epoch from checkpoint dictionary [OK]

Common Mistakes:

• Thinking print shows model parameters count
• Confusing optimizer state with epoch
• Assuming missing keys cause error here

4. You tried to resume training but got an error: RuntimeError: Error(s) in loading state_dict. What is the most likely cause related to checkpointing?

medium

A. The training data was modified after checkpointing

B. The checkpoint file was saved with torch.load instead of torch.save

C. The model architecture changed after saving the checkpoint

D. The optimizer state was not saved in the checkpoint

Solution

1. Step 1: Understand error meaning

   Loading state_dict errors usually happen if model layers differ from saved checkpoint.

2. Step 2: Connect error to checkpoint cause

   If model architecture changed after saving, weights won't match, causing this error.

3. Final Answer:

   The model architecture changed after saving the checkpoint -> Option C

4. Quick Check:

   State_dict error = architecture mismatch [OK]

Hint: Mismatch model layers cause state_dict loading errors [OK]

Common Mistakes:

• Confusing save/load functions causing error
• Assuming missing optimizer state causes this error
• Blaming training data changes for state_dict error

5. You want to checkpoint your training every 5 epochs to avoid losing progress. Which approach best preserves training progress including optimizer state and epoch count?

hard

A. Save a dictionary with model.state_dict(), optimizer.state_dict(), and current epoch number

B. Save only model.state_dict() every 5 epochs

C. Save optimizer.state_dict() and epoch number but not model weights

D. Save the training data batch every 5 epochs

Solution

1. Step 1: Identify what preserves full training state

   Saving model weights, optimizer state, and epoch number allows full resume.

2. Step 2: Compare options

   Only saving model weights misses optimizer info; saving optimizer and epoch without model is incomplete; saving data batch doesn't preserve progress.

3. Final Answer:

   Save a dictionary with model.state_dict(), optimizer.state_dict(), and current epoch number -> Option A

4. Quick Check:

   Checkpoint = model + optimizer + epoch [OK]

Hint: Checkpoint all: model, optimizer, and epoch for full resume [OK]

Common Mistakes:

• Saving only model weights loses optimizer progress
• Ignoring epoch number causes restart from zero
• Saving training data batch does not preserve model state