Saving Model Weights in PyTorch

The recommended way to save a PyTorch model is to save its state_dict, passed as the first argument to torch.save.

Programming

Saving a Model in One Call

Saving a trained PyTorch model is one of the most common operations you will perform, and it comes down to a single function call. Instead of starting with a definition, this lesson starts with the code. The save call you will examine has two arguments: an object produced from the model, and a filename. Keeping those two arguments separate in your mind is the key to the rest of the lesson.

What do you think happens?

The save call has an object in its first position. Which object do you think it is: the model itself, the output of model.state_dict(), or the filename?

The model object itself
The output of model.state_dict()
The filename string

Reveal answer

Answer:

The output of model.state_dict()

The recommended way to save a PyTorch model is to save its state dictionary, passed as the first argument to torch.save. The first argument is therefore model.state_dict(), not model, and the second argument is the destination filename.

Tracing the Save Call Step by Step

python

torch.save(model.state_dict(), "model.pth")

torch.save(model.state_dict(), "model.pth")

Walk through it

1
Python first evaluates model.state_dict(), which asks the model for its state dictionary.
2
The resulting dictionary is passed as the first argument to torch.save.
3
The string "model.pth" is passed as the second argument and names the destination file.
4
torch.save writes the dictionary to a file with that name.

Notes on the code

model.state_dict() is the first argument: the object being saved
"model.pth" is the second argument: the destination filename

The trace runs in a fixed order, and 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. A save call has only a few moving parts, and each one can be checked on its own. The diagram below lays out those steps in order.

At which step in the sequence would a wrong result first show up?

Read the flowchart as a checklist. If the saved file does not contain what you expected, you can check each step in turn: was the model in the state you intended, did state_dict() return the layer-to-parameter pairs, did torch.save receive the object you meant it to receive, and did the file end up under the name you gave. The first step where the observed result departs from the expected one is where the problem first appears.

What a State Dictionary Holds

Definition

A state dictionary is a dictionary that maps each layer of a model to its parameters. The saved object is therefore organized by layer.

The state dictionary is a snapshot of the learned values of the model, arranged by layer. When you save it, you are saving those layer-to-parameter pairs. The source identifies this as the recommended form for saving a PyTorch model. Because the first argument in the trace is model.state_dict() and not model, the object written to the file is this organized snapshot rather than the model object as a whole.

What does the state dictionary contain, and how does each layer key point to its parameters?

Read the diagram from the top. The state dictionary is one object. Each entry inside it is keyed by a layer, and the value attached to that key is that layer's parameters. The saved file therefore keeps the association between each layer and its own values, which is exactly what the source means by organizing the object by layer.

Naming the File and the .pth Convention

In the save call, the string "model.pth" names the file that stores the state dictionary. This second argument is the destination filename, and it is separate from the content being saved. Changing the filename changes where the data is written; it does not change the state dictionary itself.

Following the Data From Model to Disk

The data path has three stages. The model produces its state dictionary. That state dictionary is the first argument to torch.save. The filename is the destination, and it sits at the end of the path without altering the object in the middle of it. The diagram traces that path from the model to the file on disk.

Which object moves into torch.save, what is extracted from the model along the way, and where does the data end up on disk?

Mistakes When Saving Model Weights

Common mistake
Passing the model object itself as the first argument instead of its state dictionary
What it looks like
```
torch.save(model, "model.pth")
```
Why it fails
The source identifies the state dictionary as the recommended form for saving a PyTorch model, and the first argument in the trace is model.state_dict() and not model.
Do this instead
```
torch.save(model.state_dict(), "model.pth")
```
Common mistake
Believing the filename changes what is saved
Why it fails
The second argument is the destination filename, separate from the content being saved. It does not alter the object in the first position.
Do this instead
Check the first argument to see what content is being written, and check the second argument to see where it goes.
Common mistake
Assuming the .pth extension is required
Why it fails
The source states that .pth is a convention, not a requirement, and any file extension could technically be used.
Do this instead
Use .pth as the PyTorch convention, and understand that the extension is a naming choice rather than a rule enforced by the save call.
Common mistake
Debugging a save call by inspecting only the final file instead of checking each step in order

Practice Tracing a Save Call

MEDIUM

A colleague saves a model with the call torch.save(model.state_dict(), "checkpoint.pth"). When they inspect the result, the saved content does not match what they expected. Name the step in the trace (build model, call state_dict, call torch.save, write file) where the unexpected result would first appear, and explain how you reached that answer. Then say what the filename "checkpoint.pth" controls and what it does not control.

Hints

Work through the save call in the same order as the flowchart, and stop at the first step where the result could differ from expectation.
The second argument names the destination file; it does not change the object in the first argument.
The .pth extension is a convention, so the name ends the file path but does not decide the contents.

Your attempt (not saved)

Key Takeaways

The recommended way to save a PyTorch model is to pass its state dictionary, model.state_dict(), as the first argument to torch.save.
A state dictionary is a dictionary that maps 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 rather than a requirement; any file extension could technically be used.
To debug a save call, trace each step in order and find the first point where the result departs from what you expected.

Related concepts

More concepts starting with S