Save and Load Training Checkpoint
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:
None
{'saved_epoch': 3, 'weights_restored': True, 'optimizer_restored': True, 'checkpoint_keys': ['model_state_dict', 'optimizer_state_dict', 'epoch', 'loss']}- The function
checkpoint_test()starts by seeding withtorch.manual_seed(42)and creates a modelnn.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.BytesIObuffer 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
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.
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
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.