Save and Load Optimized IPEX Model#
This example illustrates how to save and load a model accelerated by IPEX.
In this example, we use a pretrained ResNet18 model. Then, by calling
InferenceOptimizer.trace(..., use_ipex=True), we can obtain a model accelerated by IPEX method. By calling
InferenceOptimizer.save(model=..., path=...) , we could save the Nano optimized model to a folder. By calling
InferenceOptimizer.load(path=..., model=original_model), we could load the IPEX optimized model from a folder.
First, prepare model. We need to load the pretrained ResNet18 model:
import torch from torchvision.models import resnet18 model_ft = resnet18(pretrained=True)
Accelerate Inference Using IPEX#
from bigdl.nano.pytorch import InferenceOptimizer ipex_model = InferenceOptimizer.trace(model_ft, use_ipex=True)
Save Optimized IPEX Model#
The saved model files will be saved at “./optimized_model_ipex” directory.
There are 2 files in optimized_model_ipex, users only need to take “ckpt.pth” file for further usage:
nano_model_meta.yml: meta information of the saved model checkpoint
ckpt.pth: pytorch state dict checkpoint for general use, describes model structure
Load the Optimized Model#
loaded_model = InferenceOptimizer.load("./optimized_model_ipex", model=model_ft)
The original model is required when we load the IPEX optimized model. As we only store the
state_dict, which is simply a python dictionary object that maps each layer to its parameter tensor, when saving the IPEX optimized model.
Note that when the model is optimized by IPEX as well as the JIT accelerator, the original model is not required when loading the optimized model.
Inference with the Loaded Model#
with InferenceOptimizer.get_context(loaded_model): x = torch.rand(2, 3, 224, 224) y_hat = loaded_model(x) predictions = y_hat.argmax(dim=1) print(predictions)
📚 Related Readings