How to export a PyTorch model to ONNX

ONNX export is the handoff point between a PyTorch training or prototyping environment and runtimes that read a framework-neutral graph. It matters when a model needs to move into ONNX Runtime, an inference service, an edge target, or an inspection tool without carrying the original Python module along.

The recommended PyTorch path uses torch.onnx.export(…, dynamo=True), which captures the model through torch.export before translating supported tensor operations into an ONNX graph. Put the module in inference mode first, and use an input tuple that represents the shapes and dtypes the deployment path expects.

Saving the returned ONNXProgram is only the first checkpoint. A complete export pass also loads the file with onnx, runs onnx.checker.check_model(), and executes one CPU ONNX Runtime inference so shape problems and numeric drift are caught before the artifact leaves the development environment.

1. Install the ONNX export and runtime packages in the active PyTorch environment.

   $ python -m pip install --upgrade onnx onnxscript onnxruntime

   Run the install inside the environment that already imports the model. Install PyTorch first when the environment does not have torch.
   Related: How to install PyTorch with pip

2. Create export_onnx_smoke.py with a model, representative input, exporter call, ONNX checker, and runtime parity check.

   export_onnx_smoke.py

   import numpy as np
   import onnx
   import onnxruntime as ort
   import torch
    
    
   class SmallClassifier(torch.nn.Module):
       def __init__(self):
           super().__init__()
           self.layers = torch.nn.Sequential(
               torch.nn.Linear(4, 8),
               torch.nn.ReLU(),
               torch.nn.Linear(8, 3),
           )
    
       def forward(self, features):
           return self.layers(features)
    
    
   torch.manual_seed(7)
   model = SmallClassifier().eval()
   example_input = torch.randn(1, 4)
    
   onnx_program = torch.onnx.export(
       model,
       (example_input,),
       dynamo=True,
       input_names=["features"],
       output_names=["scores"],
   )
   onnx_program.save("small-classifier.onnx")
    
   onnx_model = onnx.load("small-classifier.onnx")
   onnx.checker.check_model(onnx_model)
    
   session = ort.InferenceSession(
       "small-classifier.onnx",
       providers=["CPUExecutionProvider"],
   )
   ort_inputs = {session.get_inputs()[0].name: example_input.numpy()}
   ort_output = session.run(None, ort_inputs)[0]
    
   with torch.no_grad():
       torch_output = model(example_input).numpy()
    
   max_difference = np.max(np.abs(torch_output - ort_output))
    
   print("exported: small-classifier.onnx")
   print(f"checked: opset={onnx_model.opset_import[0].version}")
   print(f"runtime output shape: {ort_output.shape}")
   print(f"max difference: {max_difference:.8f}")

   Replace SmallClassifier and example_input with the module and input tuple from the project. Keep model.eval() for inference exports unless the runtime intentionally needs training behavior.

3. Run the export smoke script.

   $ python export_onnx_smoke.py
   [torch.onnx] Obtain model graph for `SmallClassifier([...]` with `torch.export.export(..., strict=False)`...
   ##### snipped #####
   exported: small-classifier.onnx
   checked: opset=20
   runtime output shape: (1, 3)
   max difference: 0.00000001

   Fix exporter errors before handing the file to a runtime. Unsupported Python control flow, custom operators, or shape assumptions can produce a failed export or an artifact that the checker and runtime smoke test reject.

4. Remove the smoke artifacts if they were created only for validation.

   $ rm export_onnx_smoke.py small-classifier.onnx

   Keep the ONNX file instead when it is the real deployment artifact, and store it under the project's model artifact path with the same input and output names that downstream runtime code will use.