When we using PyTorch, a useful Python deep learning framework for model training, maybe sometimes we forget to “store” the trained model, even unaware of that. (Such as me in the past)
It is a very important to save the model and load it when you want to use it.
For more detailed actual combat, you can actually refer to what I wrote before: [PyTorch] Build a GAN model to generate false MNIST pictures
Of course you can also refer the official tutorial: https://pytorch.org/tutorials/beginner/saving_loading_models.html
How to save the model
It is so easy to save a trained model:
torch.save(Model, 'Save_File_Name.pth')
Model is the object variable name of our trained model.
It is worth noting that this method stores the entire model, not just the weights, so after we loading this model in, we have no need to redefined the model architecture.
How to load the model we saved
Just like what I just said, we just need to load it directly:
torch.save(Model, 'Save_File_Name.pth')
Then we can directly start using the model. However, if you do not intend to continue training the model, but only instead to use it, remember to set the model to evaluation mode.
model.eval()
And put the test data into the block of the model and place it under the code below:
with torch.no_grad():
In this way, you will not continue to train.