torch.onnx.export converts a PyTorch model into an ONNX file by running it with a representative input and capturing the computations it performs. The input values can be arbitrary, but its shape and type must be correct. section 17.2.1
Export a model
Provide the model, a dummy tensor, and an output path. Setting dynamo=True enables the Dynamo-based export path. section 17.2.1 section 17.2.2
onnx_model_path = "simple_model.onnx"
dummy_input = torch.randn(1, 10)
torch.onnx.export( #1
model,
dummy_input,
onnx_model_path
)
This listing shows the minimal export call for a model receiving a batch of ten-element inputs. section 17.2.1
The resulting ONNX file can be executed with ONNX Runtime, compiled for an edge device, or uploaded to a cloud service that exposes it through a REST endpoint. section 17.2.1
import onnxruntime
sess = onnxruntime.InferenceSession("simple_model.onnx") #1
batch = dummy_input.numpy()
input_name = sess.get_inputs()[0].name
pred_onnx, = sess.run(None, {input_name: batch})
This listing loads the exported file with ONNX Runtime and runs inference using a NumPy batch. section 17.2.1
A practical export workflow is:
- Keep the source PyTorch model so changes can be exported again. section 17.2.2
- Export the computation graph, then pass it to a compiler for target-specific low-level code. section 17.2.2
