View the runnable example on GitHub

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

[ ]:

InferenceOptimizer.save(ipex_model, "./optimized_model_ipex")

Load the Optimized Model#

[ ]:

loaded_model = InferenceOptimizer.load("./optimized_model_ipex", model=model_ft)

📝 Note

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