Converting a PyTorch Model¶
This page provides instructions on how to convert a model from the PyTorch format to the OpenVINO IR format using Model Optimizer.
Model Optimizer Python API allows the conversion of PyTorch models using the
(Experimental) Converting a PyTorch model with PyTorch Frontend¶
Example of PyTorch model conversion:
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
example_input parameters to be set.
example_inputis used as example input for model tracing.
input_shapeis used for constructing a float zero-filled torch.Tensor for model tracing.
Example of using
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.zeros(1, 3, 100, 100))
example_input accepts the following formats:
tuplewith tensors (
dictionarywhere key is the input name, value is the tensor (
convert_model() method returns
ov.model that you can optimize, compile, or serialize into a file for subsequent use.
Exporting a PyTorch Model to ONNX Format¶
Currently, the most robust method of converting PyTorch models is exporting a PyTorch model to ONNX and then converting it to IR. To convert and deploy a PyTorch model, 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')
As of version 1.8.1, not all PyTorch operations can be exported to ONNX opset 9 which is used by default.
It is recommended to export models to opset 11 or higher when export to default opset 9 is not working. In that case, use
opset_version option of the
torch.onnx.export. For more information about ONNX opset, refer to the Operator Schemas page.