-1

I am new to Deep learning and I want to know, how can I save the final model in Pytorch? I tried some things that were mentioned but I got confused with, how to save the model and how to load it back?

desertnaut
  • 57,590
  • 26
  • 140
  • 166
asif
  • 25
  • 1
  • 9
  • Please notice that Pytorch questions are expected to be tagged as such (edited). – desertnaut Mar 26 '21 at 22:50
  • 1
    Does this answer your question? [Best way to save a trained model in PyTorch?](https://stackoverflow.com/questions/42703500/best-way-to-save-a-trained-model-in-pytorch) – desertnaut Mar 26 '21 at 22:52

1 Answers1

2

to save:

# save the weights of the model to a .pt file
torch.save(model.state_dict(), "your_model_path.pt")

to load:

# load your model architecture/module
model = YourModel()
# fill your architecture with the trained weights
model.load_state_dict(torch.load("your_model_path.pt"))
Theodor Peifer
  • 3,097
  • 4
  • 17
  • 30
  • I save the model using a torch.save() method, but I have a problem now understanding how I will load it. Do I have to create a different program for that and if yes, which parameters I have to pass. how many things will the load function take from the saved model.? – asif Mar 27 '21 at 04:36
  • what do you mean by "different program"? You can load the .pt file from any script you want. I normally have two functions, train() and test(). In train() I train the model and in every epoch I save it. In test() I just load the model.. – Theodor Peifer Mar 27 '21 at 09:35
  • can we put the save function and the load function in the same program, when we have to run it later? – asif Mar 27 '21 at 12:09