Storing and loading models

Once training is complete, we'll want to store the trained parameters in a file for deployment and future use.

 

There are two ways of doing so.

 

Let's look at the first method:

 

1. First, we will store the model parameters or state_dict in a file:

# define path2weights
path2weights="./models/weights.pt"

# store state_dict to file
torch.save(model.state_dict(), path2weights)

2. To load the model parameters from the file, we will define an object of the Net class:

# define model: weights are randomly inintiated
_model = Net()

3. Then, we will load state_dict from the file:

weights=torch.load(path2weights)

4. Next, we will set state_dict to the model:

_model.load_state_dict(weights)

In the first method, we stored state_dict or model parameters only.

 

Whenever we need the trained model for deployment, we have to create an object of the model, then load the parameters from the file, and then set the parameters to the model.

 

This is the recommended method by PyTorch creators.

 

Let's look at the second method:

 

1. First, we will store the model in a file:

# define a path2model
path2model="./models/model.pt"

# store model and weights into a file
torch.save(model, path2model)

2. To load the model parameters from the file, we will define an object of the Net class:

#define model: weights are randomly initiated
_model = Net()

3. Then, we will load the model from the local file:

_model=torch.load(path2model)

In second method, we stored the model into a file.

 

In other words, we stored both the model and state_dict into one file. 

 

Whenever we need the trained model for depolyment, we need to create an object of the Net class. 

 

Then, we loaded the model from the file. 

 

So, there is no actual benefit of doing this compared to the previous method.

 

※ reference: pytorch computer vision code book