Saving and Loading PyTorch Models
torch.save() writes data to a file. In the source example, the data is model.state_dict() and the file is "model.pth".
Why Trained Parameters Need a File
After a model has been trained, its learned parameters are the result of all that work. If those parameters exist only in the running Python session, they are lost when the session ends. Saving them to a file lets you keep that result and reuse it later. This article examines one line of PyTorch code that does exactly this job, and it works through what happens at each step of that line.
Before reading any further, run the line in your head as if it were the last line of a training script, where model is a trained PyTorch model. Ask yourself two things: what is being written, and where it ends up. Write down your prediction before you read the explanation. Comparing your prediction with what the code actually does is the fastest way to find gaps in your mental model.
What the State Dictionary Holds
model.state_dict() returns a dictionary containing the state of the model, including its parameters. It does not write a file on its own. The returned dictionary is a snapshot of the model's state, represented as a plain Python dictionary.
The key idea is that the dictionary is a plain collection of the model's parameters. Because it is a dictionary of parameters rather than the model object itself, what you end up holding is the learned values in a simple container. A dictionary is organized as keys mapped to values, so each entry in a state dictionary pairs a name with the parameter it stands for. The diagram below shows the shape of that mapping using illustrative entries; a real model's entries depend on the layers it was built with.
Three Roles in One Save Line
- torch.save: the function that writes data to a file
model.state_dict(): the data to be written, a dictionary of the model's state- "model.pth": the destination, the name of the file
The line contains exactly three roles. The first is the data, which is the state dictionary produced by calling state_dict() on the model. The second is the function that performs the write, which is torch.save. The third is the destination, the string "model.pth", which names the file. Keep these three roles separate in your head, because most save failures come from confusing one role with another.
Tracing the Data Path
The data moves in a fixed order. The model holds its parameters. The call model.state_dict() reads those parameters and returns them as a dictionary. That dictionary is passed as the first argument to torch.save(), which writes it to the file named by the second argument. The diagram below traces this path from left to right.
Tracing one save from start to finish
A training script ends with torch.save(model.state_dict(), "model.pth"), where model is trained. What is written, by which function, and to what destination?
- 1Step 1: the model holds the parameters
The trained model object contains its learned parameters. The model itself is not the thing being written.
- 2Step 2: state_dict() produces the dictionary
The call
model.state_dict()returns a dictionary containing the state of the model, including its parameters. Nothing is written to disk at this point. - 3Step 3: torch.save() takes the dictionary
The dictionary is the first argument to
torch.save(), so it is the data to be written. - 4Step 4: the file receives the data
The second argument, "model.pth", is the destination. The dictionary of parameters is written to that file.
The file model.pth stores the dictionary of learned parameters produced by model.state_dict(). It does not store the model object itself.
Snapshot Versus Whole Model
Because the snapshot is a plain dictionary of the model's parameters, the saved file holds the learned values rather than the model object. The table below contrasts the two things you could hand to torch.save, so you can see what each choice puts into the file.
| Aspect | Saving model.state_dict() | Saving the model object |
|---|---|---|
| What is passed to torch.save | A plain dictionary of the model's state, including parameters | The model object itself |
| What the file stores | The learned parameter values held in the dictionary | The model object, not just its dictionary of parameters |
| Role of the model | The source the dictionary is read from | The thing being written |
| Matches this article's line | Yes: torch.save(model.state_dict(), "model.pth") | No: the line does not pass the model object |
Two possible first arguments to torch.save, and what each one puts into the file
Where the Save Can Diverge
The point where a save can produce a different result from what you intended is the moment the data passes from the model into the file. That is the first argument to torch.save. When a save does not produce the result you expect, check the three roles in order: the data, the saving function, and the destination. The flowchart shows that order.
Common mistake Passing the model object where the state dictionary belongs
What it looks liketorch.save(model, "model.pth")Why it failsThe data argument should be the dictionary produced by
model.state_dict(). The model object is not the snapshot the source describes, and the file should store the learned values.Do this insteadtorch.save(model.state_dict(), "model.pth")Common mistake Expecting state_dict() to write the file
What it looks likemodel.state_dict() # expecting model.pth to appearWhy it failsmodel.state_dict()returns a dictionary and does not write a file on its own. Nothing reaches the disk untiltorch.save()is called.Do this insteadCall torch.save() with the dictionary as its first argument.
Common mistake Writing to an unintended destination name
What it looks liketorch.save(model.state_dict(), "model.txt") when you meant model.pthWhy it failsThe second argument is the file name. Whatever string you pass is where the data goes, so a wrong name silently stores the result somewhere you may not look.
Do this insteadCheck that the second argument is exactly the file name you intended.
Predict Before You Reveal
What do you think happens?
When the line torch.save(model.state_dict(), "model.pth") runs, which call creates the file model.pth?
Reveal answer
torch.save() creates the file, because it writes data to the file named by its second argument.
model.state_dict() returns a dictionary of the model's state and does not write a file on its own. torch.save() takes that dictionary as its first argument and writes it to model.pth. The file stores the dictionary of parameters, not the model object itself.
A colleague writes torch.save(model, "checkpoint.pth") after training and says the file holds the learned weights as a state dictionary. Using the three roles, identify the data, the function, and the destination in their line. Then state where the result can diverge from the intended save and write the corrected line.
Hints
- Start with the first argument and ask whether it is
model.state_dict(). - The second argument is the destination, and here it is checkpoint.pth.
- The corrected line keeps torch.save and the destination, and changes only the data argument.
Key Takeaways
- 1
model.state_dict()returns a dictionary of the model's state, including its parameters, and does not write a file by itself. - 2
torch.save()writes its first argument to the file named by its second argument. - 3In
torch.save(model.state_dict(), "model.pth"), the three roles are the data, the saving function, and the destination. - 4A save diverges from intent most often at the data argument, so check that it is the state dictionary before checking the function and the file name.
At a glance
model.state_dict()returns a dictionary containing the model's state, including its parameters, and does not write a file on its own.torch.save()takes the data in its first argument and writes it to the file named by its second argument.- The data flows from the model to
state_dict()totorch.save()and then to the file, so each step can be checked in order. - Saving the state dictionary stores the learned values as a plain dictionary rather than the model object itself.
- When a save goes wrong, check the data argument, the saving function, and the destination file name in turn.