Loading a Saved PyTorch State Dictionary
The recommended way to save a PyTorch model is to save its state_dict, passed as the first argument to torch.save.
One Call, Two Arguments
Saving a trained PyTorch model is one of the most common operations you will perform, and it comes down to a single function call. This lesson does not begin with a definition. It begins with the code. You will watch one save call run from start to finish, and before reading any explanation, you will decide what each part of that call is doing.
The call has two arguments. The first is an object produced from the model. The second is a filename. The rest of this lesson depends on telling those two arguments apart, so keep both in mind as you read.
What do you think happens?
A save call has an object in its first position. Before you look at the trace, what do you think that object is: the model itself, or something produced from the model?
Reveal answer
It is model.state_dict(), the state dictionary produced from the model, not the model itself.
The recommended way to save a PyTorch model is to save its state_dict, passed as the first argument to torch.save. The state dictionary is a snapshot of the learned values organized by layer.
Tracing the Save Call Step by Step
Here is the save call in the form the trace describes it. The first argument is model.state_dict(), and the second argument is the string naming the destination file.
model.state_dict()runs first and produces the object that gets saved- The first argument to torch.save is that state dictionary
- "model.pth" is the second argument, the destination filename
Read the trace in the order the steps run. First, model.state_dict() produces the object that will be saved. Second, that object enters torch.save in the first position. Third, the string "model.pth" enters torch.save in the second position as the destination. Finally, the file is written. The goal is not only to see the final file. It is to be able to point to the exact step where a value could go wrong.
What a State Dictionary Holds
A state dictionary is a dictionary mapping each layer to its parameters, so the saved object is organized by layer. Saving it means saving those layer-to-parameter pairs, which the source identifies as the recommended form for saving a PyTorch model.
Think of the state dictionary as a snapshot of the learned values. Each key names a layer, and each value holds that layer's parameters. The diagram below uses a small, invented two-layer model to show the shape. The layer names fc1 and fc2 are illustrative and are not taken from the source.
Because the saved object is organized by layer, each layer's parameters stay grouped under that layer's name. This is why the first argument in the trace is model.state_dict() and not the bare model. The state dictionary is the layer-organized form of the learned values, and it is the object the source recommends saving.
Where the Filename Fits
The second argument to torch.save is the destination filename. It is separate from the content being saved. The string "model.pth" names the file that stores the state dictionary, so the filename tells the save call where to write, while the first argument tells it what to write.
Common mistake Treating the filename as part of the content being saved
Why it failsThe second argument is the destination. It is kept separate from the object in the first position.
Do this insteadRead the first argument as what is saved and the second argument as where it is saved.
Common mistake Believing the .pth extension is required
Why it failsThe source states that .pth is a convention and that any file extension could technically be used.
Do this insteadKeep .pth for consistency with PyTorch convention, but do not assume a different extension is an error by itself.
Common mistake Assuming the bare model object is what gets saved
Why it failsThe recommended form is the state dictionary, produced by
model.state_dict()and passed as the first argument.Do this insteadCheck that the first argument to torch.save is model.state_dict() when following the recommended approach.
Practice Predicting the Trace
Consider the call torch.save(model.state_dict(), "checkpoint.pth"). Trace it in order. Name the object in the first position, name the destination in the second position, and say which step you would check first if the saved file did not contain the layer-to-parameter pairs you expected.
Hints
- Start with
model.state_dict(), since it runs before torch.save receives its arguments. - The string in the second position is only the destination filename.
- The source points to the step that builds the first argument as the place where an unexpected result would first appear.
Explain, in your own words, why changing the filename from model.pth to model.bin would not change which object is saved. Then say what would have to change in the call to save a different object.
Hints
- Separate the destination argument from the content argument.
- The source says the .pth extension is a convention and not a rule.
Key Takeaways
- The recommended way to save a PyTorch model is to pass its
state_dictas the first argument to torch.save. - A state dictionary is a dictionary mapping each layer to its parameters, so the saved object is organized by layer.
- The second argument to torch.save is the destination filename, separate from the content being saved.
- The .pth extension is a PyTorch convention, not a requirement; any file extension could technically be used.
- When debugging a save call, trace the steps in order and check the step that builds the first argument first.