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)

Then we set it in evaluation mode:

[ ]:

model_ft.eval()

To enable OpenVINO acceleration for your PyTorch inference pipeline, the only change you need to made is to import BigDL-Nano InferenceOptimizer, and trace your PyTorch model to convert it into an OpenVINO accelerated module 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 with the model optimized by OpenVINO:

[ ]:

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)