Conversion between PyTorch and ONNX
In this section, we will explain how to convert a PyTorch model into an ONNX model and back again. With the conversion between TF and ONNX covered in the previous section, you should be able to convert your model between TF and PyTorch as well by the end of this section.
Converting a PyTorch model into an ONNX model
Interestingly, PyTorch has built-in support for exporting its model as an ONNX model (https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html). Given a model, all you need is the torch.onnx.export function as shown in the following code snippet:
import torch pytorch_model = ... # Input to the model dummy_input = torch.randn(..., requires_grad=True) onnx_model_path = "model.onnx" # Export the model torch.onnx.export( pytorch_model, # model being run dummy_input, ...
