Model stored in two ways pytorch

1. Save the entire network configuration information and the model parameters :

torch.save(model_object, './model.pth')

Direct loading can be used:

model = torch.load('./model.pth')

 

2. The model parameters are saved only network - recommended

torch.save(model_object.state_dict(), './params.pth')

Loading the local network will have to start import network module, and then loading parameters:

from models import AgeModel
model = AgeModel()
model.load_state_dict(torch.load('./params.pth'))