Save and Load Your PyTorch Models

Last Updated on April 8, 2023

A deep learning model is a mathematical abstraction of data, in which a lot of parameters are involved. Training these parameters can take hours, days, and even weeks but afterward, you can make use of the result to apply on new data. This is called inference in machine learning. It is important to know how we can preserve the trained model in disk and later, load it for use in inference. In this post, you will discover how to save your PyTorch models to files and load them up again to make predictions. After reading this chapter, you will know:

  • What are states and parameters in a PyTorch model
  • How to save model states
  • How to load model states

Kick-start your project with my book Deep Learning with PyTorch. It provides self-study tutorials with working code.

Let’s get started.

Save and Load Your PyTorch Models
Photo by Joseph Chan. Some rights reserved.


This post is in three parts; they are

  • Build an Example Model
  • What’s Inside a PyTorch Model
  • Accessing state_dict of a Model

Build an Example Model

Let’s start with a very simple model in PyTorch. It is a model based on the iris dataset. You will load the dataset using scikit-learn (which the targets are integer labels 0, 1, and 2) and train a neural network for this multiclass classification problem. In this model, you used log softmax as the output activation so you can combine with the negative log likelihood loss function. It is equivalent to no output activation combined with cross entropy loss function.

With such a simple model and small dataset, it shouldn’t take a long time to finish training. Afterwards, we can confirm that this model works, by evaluating it with the test set:

It prints, for example,

Want to Get Started With Deep Learning with PyTorch?

Take my free email crash course now (with sample code).

Click to sign-up and also get a free PDF Ebook version of the course.

What’s Inside a PyTorch Model

PyTorch model is an object in Python. It holds some deep learning building blocks such as various kinds of layers and activation functions. It also knows how to connect them so it can produce you an output from your input tensors. The algorithm of a model is fixed at the time you created it, however, it has trainable parameters that is supposed to be modified during training loop so the model can be more accurate.

You saw how to get the model parameters when you set up the optimizer for your training loop, namely,

The function model.parameters() give you a generator that reference to each layers’ trainable parameters in turn in the form of PyTorch tensors. Therefore, it is possible for you to make a copy of them or overwrite them, for example:

Which the result should be exactly the same as before since you essentially made the two models identical by copying the parameters.

However, this is not always the case. Some models has non-trainable parameters. One example is the batch normalization layer that is common in many convolution neural networks. What it does is to apply normalization on tensors that produced by its previous layer and pass on the normalized tensor to its next layer. It has two parameters: The mean and standard deviation, which are learned from your input data during training loop but not trainable by the optimizer. Therefore these are not part of model.parameters() but equally important.

Accessing state_dict of a Model

To access all parameters of a model, trainable or not, you can get it from state_dict() function. From the model above, this is what you can get:

The model above produces the following:

It is called state_dict because all state variables of a model are here. It is an OrderedDict object from Python’s built-in collections module. All components from a PyTorch model has a name and so as the parameters therein. The OrderedDict object allows you to map the weights back to the parameters correctly by matching their names.

This is how you should save and load the model: Fetch the model states into an OrderedDict, serialize and save it to disk. For inference, you create a model first (without training), and load the states. In Python, the native format for serialization is pickle:

You know it works because the model you didn’t train produced the same result as the one you trained.

Indeed, the recommended way is to use the PyTorch API to save and load the states, instead of using pickle manually:

The *.pth file is indeed a zip file of some pickle files created by PyTorch. It is recommended because PyTorch can store additional information in it. Note that you stored only the states but not the model. You still need to create the model using Python code and load the states into it. If you wish to store the model as well, you can pass in the entire model instead of the states:

But remember, due to the nature of Python language, doing so does not relieve you from keeping the code of the model. The newmodel object above is an instance of Multiclass class that you defined before. When you load the model from disk, Python need to know in detail how this class is defined. If you run a script with just the line torch.load(), you will see the following error message:

That’s why it is recommended to save only the state dict rather than the entire model.

Putting everything together, the following is the complete code to demonstrate how to create a model, train it, and save to disk:

And the following is how to load the model from disk and run it for inference:

Further Readings

This section provides more resources on the topic if you are looking to go deeper.


In this post, you learned how to keep a copy of your trained PyTorch model in disk and how to reuse it. In particular, you learned

  • What are parameters and states in a PyTorch model
  • How to save all necessary states from a model to disk
  • How to rebuild a working model from the saved states

Get Started on Deep Learning with PyTorch!

Deep Learning with PyTorch

Learn how to build deep learning models

...using the newly released PyTorch 2.0 library

Discover how in my new Ebook:
Deep Learning with PyTorch

It provides self-study tutorials with hundreds of working code to turn you from a novice to expert. It equips you with
tensor operation, training, evaluation, hyperparameter optimization, and much more...

Kick-start your deep learning journey with hands-on exercises

See What's Inside

4 Responses to Save and Load Your PyTorch Models

  1. Avatar
    Javier March 10, 2023 at 6:35 am #

    Thank you it was really useful, just a comment. In the final example the loading operation is shown using picklle instead of torch.load. Torch load will be more consistent with the saving code showed before.

    • Avatar
      James Carmichael March 10, 2023 at 7:59 am #

      You are very welcome Javier! Thank you for your feedback and suggestion!

  2. Avatar
    Poult August 10, 2023 at 4:48 pm #

    Hello, I think in the last windows of code, in the 31-32 line:
    it could be: ‘with open(“iris-model.pth”, “rb”) as fp:

    • Avatar
      James Carmichael August 11, 2023 at 8:15 am #

      Thank you for your feedback and suggestion Poult!

Leave a Reply