How to export a PyTorch model to ONNX

Deployment runtimes often need a model artifact that does not depend on the original Python class or training process. ONNX provides that handoff by representing the captured tensor computation as a graph that runtimes such as ONNX Runtime can load directly.

The torch.export-based exporter behind torch.onnx.export(…, dynamo=True) is the recommended PyTorch path. It captures the module with representative tensor inputs, translates supported operations, and returns an ONNXProgram that can be saved without passing a legacy output-path argument to the exporter.

An exported file still needs two separate checks before deployment. The ONNX checker catches an invalid graph, while a CPU ONNX Runtime inference and numeric comparison catch input-name, output-name, shape, and translation problems that graph validation alone cannot expose.

1. Install the ONNX exporter and CPU runtime dependencies in the active PyTorch environment.

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

   This command belongs in the environment that already imports the project model; the PyTorch package itself comes from the project's chosen CPU or accelerator package source.
   Related: How to install PyTorch with pip

2. Create export_onnx_smoke.py containing the imports and representative classifier setup.

   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)

   SmallClassifier and example_input stand in for the project module and a tensor tuple whose dtypes and dimensions match deployment traffic. The CPU smoke test expects the module and inputs on CPU.

3. Append the ONNXProgram export and serialization block to export_onnx_smoke.py.

   onnx_program = torch.onnx.export(
       model,
       (example_input,),
       dynamo=True,
       input_names=["features"],
       output_names=["scores"],
   )
   onnx_program.save("small-classifier.onnx")

   The example exports a fixed input shape. Supply dynamic_shapes to torch.onnx.export() when a deployment dimension such as batch size must vary.

4. Append the graph checker and CPU ONNX Runtime inference block to export_onnx_smoke.py.

   onnx_model = onnx.load("small-classifier.onnx")
   onnx.checker.check_model(onnx_model)
    
   session = ort.InferenceSession(
       "small-classifier.onnx",
       providers=["CPUExecutionProvider"],
   )
   ort_output = session.run(
       ["scores"],
       {"features": example_input.numpy()},
   )[0]

5. Append the PyTorch parity assertion and result reporting block to export_onnx_smoke.py.

   with torch.no_grad():
       torch_output = model(example_input).numpy()
    
   np.testing.assert_allclose(torch_output, ort_output, rtol=1e-5, atol=1e-6)
   max_difference = np.max(np.abs(torch_output - ort_output))
    
   print("exported: small-classifier.onnx")
   print(f"runtime input: {session.get_inputs()[0].name}")
   print(f"runtime output: {session.get_outputs()[0].name}")
   print(f"runtime output shape: {ort_output.shape}")
   print(f"max difference: {max_difference:.8f}")

   assert_allclose() raises an error when the runtime output exceeds the chosen tolerances, so the printed summary appears only after both implementations agree.

6. Run export_onnx_smoke.py for the complete ONNX export smoke test.

   $ python export_onnx_smoke.py
   ##### snipped #####
   exported: small-classifier.onnx
   runtime input: features
   runtime output: scores
   runtime output shape: (1, 3)
   max difference: 0.00000001

   Any exporter, checker, runtime, or parity error blocks the handoff. Unsupported operations, Python control flow, mismatched names, and shape constraints require correction before the ONNX file is suitable for a downstream runtime.