I’m trying to figure out what’s the best way to save a model trained with Pytorch and load it for inference, and I was wondering about the different possible approaches.
Let’s say I successfully train a model, as far as I understand I can use:
- Complete Model Saving:
# save the model
torch.save(model, saved_model_path)
# load model directly with
loaded_model = torch.load(saved_model_path)
# use it for inference
output = loaded_model(input)
- State Dict Saving:
# save only the state_dict after training
torch.save(model.state_dict(), saved_model_path)
# need to create an instance of the model with the same architecture and then load the parameters using
model = SomeModelConstructor()
model.load_state_dict(state_dict)
# use it for inference
output = loaded_model(input)
And I seem to understand saving only the state_dict has the advantage of reducing the file size, while the con is that I have to recreate the model instance. Is that correct?
Also, I see I can also export the model to torchscript and load it:
# export to torchscript and save locally
model_scripted = torch.jit.script(model)
torch.jit.save(model_scripted, scripted_model_path)
# load the scripted model
loaded_scripted_model = torch.jit.load(scripted_model_path)
# use for evaluation
loaded_scripted_model.eval()
with torch.no_grad():
output = loaded_scripted_model(input)
Assuming I’m exporting the model to script from Python and loading it back always from Python, are there any advantages in using the scripted model instead of the saved one (for example, inference speed because the scripted version is better optimized or something similar?)
