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 bookDeep Learning with PyTorch. It providesself-study tutorials withworking code.
Let’s get started.
Save and Load Your PyTorch Models Photo byJoseph Chan. Some rights reserved.
Overview
This post is in three parts; they are
Build an Example Model
What’s Inside a PyTorch Model
Accessingstate_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:
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,
1
optimizer=optim.Adam(model.parameters(),lr=0.001)
The functionmodel.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:
1
2
3
4
5
6
7
8
9
10
# create a new model
newmodel=Multiclass()
# ask PyTorch to ignore autograd on update and overwrite parameters
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 hasnon-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 ofmodel.parameters() but equally important.
Accessingstate_dict of a Model
To access all parameters of a model, trainable or not, you can get it fromstate_dict() function. From the model above, this is what you can get:
It is calledstate_dict because all state variables of a model are here. It is anOrderedDict object from Python’s built-incollections module. All components from a PyTorch model has a name and so as the parameters therein. TheOrderedDict 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 anOrderedDict, 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:
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. Thenewmodel object above is an instance ofMulticlass 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 linetorch.load(), you will see the following error message:
It providesself-study tutorials withhundreds 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
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.
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: model.load_state_dict(torch.load(fp))’
