Save and Load Training Checkpoint

hardPROModel SerializationPyTorch Advanced

Problem Statement

Create a complete training checkpoint with model, optimizer, and training state, save to a buffer, and restore from it.

Background

Training checkpoints include not just model weights but also optimizer state (momentum buffers, learning rates) and training metadata (epoch, loss). Use io.BytesIO as an in-memory buffer.

Your Task

The starter code trains for 3 steps and prepares a checkpoint dict. Save the complete training checkpoint (model state, optimizer state, epoch, loss) to an in-memory buffer, then restore everything from the buffer into fresh model and optimizer instances.

Output Format

Returns a dictionary with "saved_epoch", "weights_restored", "optimizer_restored", and "checkpoint_keys".

Example

Input

None

Output

{'saved_epoch': 3, 'weights_restored': True, 'optimizer_restored': True, 'checkpoint_keys': ['model_state_dict', 'optimizer_state_dict', 'epoch', 'loss']}

Reasoning

The function checkpoint_test() starts by seeding with torch.manual_seed(42) and creates a model nn.Linear(3, 2) with an Adam optimizer, then runs 3 training steps with input [[1.0, 2.0, 3.0]] and target [[1.0, 0.0]] to calculate the loss.
After training, it saves a checkpoint to an io.BytesIO buffer containing the model state dictionary, optimizer state dictionary, epoch (3), and the last loss value.
A new model and optimizer are created with the same architecture but a fresh seed (99), and the checkpoint is loaded from the buffer to restore the model and optimizer states.
The function then verifies that the model weights match after loading and checks if the optimizer state was loaded, resulting in the output dictionary with the saved epoch, weights restoration status, optimizer restoration status, and checkpoint keys.

Constraints

Use io.BytesIO (no file I/O)
Save model + optimizer + metadata
Verify full restoration

Learn this in

PyTorch FundamentalsFoundations CI/CD & PipelinesSystems & Networking

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

Introduction to Model Serialization

Model serialization is the process of converting a model's state into a format that can be written to a file or stored in a database. This is useful for saving and loading trained models, allowing us to resume training from a previous checkpoint or deploy the model in a different environment. In PyTorch, model serialization is typically done using the torch.save() and torch.load() functions.

Understanding Training Checkpoints

A training checkpoint is a snapshot of the training process at a particular point in time. It typically includes the model's weights, optimizer state, and other relevant metadata such as the current epoch and loss value. Saving checkpoints is essential for long-running training processes, as it allows us to recover from failures or interruptions. In this problem, we need to save a checkpoint that includes the model state dictionary, optimizer state dictionary, epoch, and loss value.

PyTorch Serialization APIs

PyTorch provides several APIs for serializing and deserializing models and checkpoints. The torch.save() function is used to save a PyTorch object (such as a model or optimizer) to a file, while the torch.load() function is used to load a saved object. The state_dict attribute of a PyTorch model or optimizer returns a dictionary representing its current state, which can be saved and loaded using torch.save() and torch.load(). In this problem, we will use io.BytesIO as an in-memory buffer to store the checkpoint.

Algorithm/Approach

The general approach to solving this problem involves the following steps:

Create a model and optimizer, and train the model for a few steps to generate a checkpoint.
Save the checkpoint to a buffer using torch.save().
Create a new model and optimizer, and load the checkpoint from the buffer using torch.load().
Verify that the weights and optimizer state have been restored correctly.

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