How to save the best model?

2

Hi, Talita:
You can save your model by either of the following methods.

# Method 1
torch.save(model, 'best-model.pt') 
# Method 2
torch.save(model.state_dict(), 'best-model-parameters.pt') # official recommended

The difference between two methods is that the first one saves the whole model which includes project-specific classes and your best parameters, while the second one just saves your best parameters.:smiley:

By the way, you can load the model by these codes.

the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

References:

  1. recommend-saving-models
6 Likes