Converting a PyTorch Model¶
This page provides instructions on how to convert a model from the PyTorch format to the OpenVINO IR format.
The conversion is a required step to run inference using OpenVINO API. It is not required if you choose to work with OpenVINO under the PyTorch framework, using its torch.compile feature.
Converting a PyTorch model with PyTorch Frontend¶
To convert a PyTorch model to the OpenVINO IR format, use the OVC API (superseding the previously used tool, MO). To do so, use the
convert_model() method, like so:
import torchvision import torch from openvino.tools.mo import convert_model model = torchvision.models.resnet50(pretrained=True) ov_model = convert_model(model)
Following PyTorch model formats are supported:
Converting certain PyTorch models may require model tracing, which needs the
parameter to be set, for example:
import torchvision import torch from openvino.tools.mo import convert_model model = torchvision.models.resnet50(pretrained=True) ov_model = convert_model(model, example_input=torch.randn(1, 3, 100, 100))
example_input accepts the following formats:
tuplewith tensors (
dictionarywhere key is the input name, value is the tensor (
convert_model will produce inputs of the model with dynamic rank or dynamic type.
Such model may not be supported by the hardware chosen for inference. To avoid this issue,
input argument of
convert_model. For more information, refer to Convert Models Represented as Python Objects.
convert_model() method returns
ov.Model that you can optimize, compile, or save to a file for subsequent use.
Exporting a PyTorch Model to ONNX Format¶
It is also possible to export a PyTorch model to ONNX and then convert it to OpenVINO IR. To convert and deploy a PyTorch model this way, follow these steps:
PyTorch models are defined in Python. To export them, use the
torch.onnx.export() method. The code to
evaluate or test the model is usually provided with its code and can be used for its initialization and export.
The export to ONNX is crucial for this process, but it is covered by PyTorch framework, therefore, It will not be covered here in detail.
For more information, refer to the Exporting PyTorch models to ONNX format guide.
To export a PyTorch model, you need to obtain the model as an instance of
torch.nn.Module class and call the
import torch # Instantiate your model. This is just a regular PyTorch model that will be exported in the following steps. model = SomeModel() # Evaluate the model to switch some operations from training mode to inference. model.eval() # Create dummy input for the model. It will be used to run the model inside export function. dummy_input = torch.randn(1, 3, 224, 224) # Call the export function torch.onnx.export(model, (dummy_input, ), 'model.onnx')