PyTorch Model Persistence

1. The Context: Saving Your Work

In any software project, you save your code. In deep learning, you need to save your trained model. This is the final step in the typical development workflow before you might deploy or share your model.

2. The Core Concept: The state_dict

Before we start saving things, we need to understand what we're saving. The most important component of a trained model is its set of learned parameters (weights and biases). In PyTorch, these are neatly contained within an object called the state_dict.

Think of the state_dict as the "brain" or the "DNA" of your model. It contains all the learned knowledge, but not the model's architecture (the "body"). This distinction is key to understanding the different ways to save a model.

3. Two Methods for Saving and Loading Models

PyTorch offers two main approaches for saving a model. One is flexible and recommended for production and sharing, while the other is simpler but more brittle.

Method 1: Saving the state_dict (Recommended)

This approach saves only the model's parameters. It is the most flexible and recommended method because it decouples the learned weights from the specific Python code that defined the model's architecture.

To save:
You simply save the state_dict dictionary using torch.save().

# Assuming 'model' is your trained torch.nn.Module instance
torch.save(model.state_dict(), 'model_weights.pth')

To load:
This is a two-step process:

1. First, you must create an instance of your model architecture.
2. Then, you load the saved parameters from the state_dict into this new model instance.

# Create an instance of the model architecture first
model = YourModelClass(*args, **kwargs)




# Load the state dictionary
model.load_state_dict(torch.load('model_weights.pth'))




# **Crucially**, set the model to evaluation mode for inference
model.eval()

The model.eval() call is vital. It tells PyTorch to turn off layers like Dropout and Batch Normalization that behave differently during training and inference.

Method 2: Saving the Entire Model (The "Lazy" Way)

This method saves the entire model object using Python's pickle utility. While it seems simpler, it has a major drawback.

To save:

torch.save(model, 'entire_model.pth')

To load:

model = torch.load('entire_model.pth')
model.eval()

The problem is that this method binds the saved file to the exact class definitions and directory structure used during saving. If you try to load this model in a different project or after refactoring your code, it will likely break. Therefore, this method is not recommended for any serious work.

Let's watch a video that demonstrates both of these methods with code.

4. Checkpointing: Saving Your Progress to Resume Training

Training a deep learning model can take hours or even days. What if the power goes out, or your program crashes? If you only save the final model, you lose all progress. This is where checkpointing comes in.

A checkpoint is a snapshot of the entire training state, saved periodically (e.g., after every epoch). It includes not just the model's weights, but everything needed to resume training exactly where you left off.

A typical checkpoint dictionary includes:

• The current epoch number
• The model's state_dict
• The optimizer's state_dict (this is important, as optimizers like Adam have internal states)
• The latest training loss
• Any other relevant information

Now, let's watch a detailed demonstration of how to implement this in a training script.

A common strategy is to save a checkpoint at the end of each epoch and also to keep track of the best-performing model on a validation set, overwriting a best_model.pth file whenever a new best is found.

5. A Critical Detail: Saving and Loading Across Devices (CPU/GPU)

A frequent source of errors occurs when you train a model on a GPU but then try to load it on a machine with only a CPU (or vice-versa). PyTorch requires you to be explicit about this device mapping.

The key is the map_location argument in the torch.load() function.

Here's a quick summary of the most common case: loading a GPU-trained model onto a CPU.

# Specify map_location when loading
checkpoint = torch.load(PATH, map_location=torch.device('cpu'))




# Now load the state_dict as usual
model.load_state_dict(checkpoint['model_state_dict'])

Failing to set map_location when loading a GPU-trained model on a CPU-only machine will result in a runtime error. This is a perfect lead-in to our next topic.

Conclusion

You now have the practical skills to manage the lifecycle of your PyTorch models. This is a fundamental part of the deep learning engineering process, enabling you to save, share, and reuse your work effectively.

Key Takeaways:

• The state_dict is a Python dictionary containing a model's learnable parameters and is the key to flexible model saving.
• The recommended method for saving is to store only the model.state_dict(), which decouples the weights from the code. Loading requires you to first instantiate the model architecture.
• Checkpoints are crucial for long training jobs. They are dictionaries containing the model state, optimizer state, epoch, and loss, allowing you to resume training seamlessly.
• When loading a model on a different device than it was trained on (e.g., GPU to CPU), you must use the map_location argument in torch.load() to prevent errors.

Preview of the Next Lesson:
We've just seen how critical it is to manage models between CPU and GPU. In the next and final lesson of this module, "Use GPU acceleration for training deep learning models," we will dive into how to leverage GPUs to dramatically reduce training time and tackle larger, more complex models.