Simple way to save and load model in pytorch

Sanpreet Singh
3 min readDec 25, 2019
Simple way to save and load model in pytorch
Simple way to save and load model in pytorch

Frameworks such as tensorflow, caffe, pytorch or keras are hot burning topics these days. Because of these, artificial intelligence has gained significant importance and fast development. Just imaging in 19th century, operations like object detection, image segmentation, image classification seems impossible and now these are common. A lot of videos has been published day by day to demonstrate these. Computers are being replaced by Graphical Processing Units and era of GPU’s has begun. This particular post is about saving and loading model in pytorch. I am writing these articles because there may be many passionate coders who are starting their career in this field, so it is vital for them to know how to load model after training has been completed so that much of their time and resources could be saved. I had also written similar articles in the past with respect to tensorflow and keras. If you wish to read them, then below are the links for them.

Loading saved keras model and continue training from last epoch

Save and Restore Tensorflow Model

Moreover, I also use to write on quora covering such worthful topics which may also help you whose link I am sharing below. If you wish you may read from there also.

https://www.quora.com/profile/Sanpreet-Singh-10

Just coming back to the article , “Simple way to save and load model in pytorch” I tried to make it simple by uploading the code on GitHub so that readers can see the things practically. Along with the code, well documented readme file is also written which will help you in a better way. Please find the link for github from below

https://github.com/sanpreet/Simple-way-to-save-and-load-model-in-pytorch

Along with the above, if readers would like to see some ideas from the images instead of content then please visit my Pinterest from below link where some creative ideas are being shared

https://in.pinterest.com/sanpreetsingh33/

Simple way to save and load model in pytorch
Simple way to save and load model in pytorch

Saving Model in pytorch

Pytorch Model is saved either in .pt or .pth format. To save the model, I defined checkpoint directory where the model is saved. torch.save will save the model. In the code provided in the GitHub link simple way to save and load model in pytorch , I have used a dictionary where both model state and optimizer is saved. Please have a look at below screenshot.

Saving Model in Pytorch
Saving Model in Pytorch

Loading Model in pytorch

Since, one has saved model in the previous step, one can load the model using torch.load(checkpoint_file). Process is very simple. Loading the model from checkpoint directory saves a lot of time because one does not need to start the training again to give prediction. Code is already mentioned in the github but still for better understanding I am adding the snippet here.

Loading Model in Pytorch
Loading Model in Pytorch

Let us also see the results after loading the model

Results after loading Model in Pytorch
Results after loading Model in Pytorch

Conclusion

In this short article, we have together learned how to save and load model in pytorch. Github code with well documented readme file is also included with this article. Snippets included in this article helps readers to visualize the code in a better way. At the end, I would like to say that there are a lot of things to be learned and with more practice and understanding it can be achieved. Please keep reading useful stuff and starts practice coding because there is no other way out.

Originally published at http://ersanpreet.wordpress.com on December 25, 2019.

--

--