View the runnable example on GitHub

Accelerate PyTorch Inference using OpenVINO#

You can use InferenceOptimizer.trace(..., accelerator='openvino') API to enable the OpenVINO acceleration for PyTorch inference. It only takes a few lines.

Let’s take an ResNet-18 model pretrained on ImageNet dataset as an example. First, we load the model:

[ ]:

from torchvision.models import resnet18

model_ft = resnet18(pretrained=True)

To enable OpenVINO acceleration for your PyTorch inference pipeline, the major change you need to make is to import BigDL-Nano InferenceOptimizer, and trace your PyTorch model to convert it into an OpenVINO accelerated model for inference:

[ ]:

import torch
from bigdl.nano.pytorch import InferenceOptimizer

ov_model = InferenceOptimizer.trace(model_ft,
                         accelerator="openvino",
                         input_sample=torch.rand(1, 3, 224, 224))

📝 Note

input_sample is the parameter for OpenVINO accelerator to know the shape of the model input. So both the batch size and the specific values are not important to input_sample. If we want our test dataset to consist of images with \(224 \times 224\) pixels, we could use torch.rand(1, 3, 224, 224) for input_sample here.

Please refer to API documentation for more information on InferenceOptimizer.trace.

You could then do the normal inference steps under the context manager provided by Nano, with the model optimized by OpenVINO:

[ ]:

with InferenceOptimizer.get_context(ov_model):
    x = torch.rand(2, 3, 224, 224)
    # use the optimized model here
    y_hat = ov_model(x)
    predictions = y_hat.argmax(dim=1)
    print(predictions)

📝 Note

For all Nano optimized models by InferenceOptimizer.trace, you need to wrap the inference steps with an automatic context manager InferenceOptimizer.get_context(model=...) provided by Nano. You could refer to here for more detailed usage of the context manager.