How to Export and Convert Models to ONNX or TorchScript in GitHub Actions

Export the model, run onnx.checker or a TorchScript reload, then compare outputs against the source model within a tolerance so a bad conversion fails CI.

Convert the model to ONNX with torch.onnx.export (or trace to TorchScript), validate the graph with onnx.checker.check_model, and assert the exported model matches the original outputs within a numerical tolerance.

Steps

• Export the model to ONNX or TorchScript.
• Validate the graph (onnx.checker or a TorchScript reload).
• Compare outputs to the source model within a tolerance.

Workflow

steps:
  - run: pip install torch onnx onnxruntime numpy
  - name: Export and verify ONNX
    run: |
      python - <<'PY'
      import torch, onnx, onnxruntime as ort, numpy as np
      from model import build
      m = build().eval()
      x = torch.randn(1, 3, 224, 224)
      torch.onnx.export(m, x, "model.onnx", opset_version=17,
                        input_names=["input"], output_names=["output"])
      onnx.checker.check_model(onnx.load("model.onnx"))
      ref = m(x).detach().numpy()
      got = ort.InferenceSession("model.onnx").run(None, {"input": x.numpy()})[0]
      np.testing.assert_allclose(ref, got, rtol=1e-3, atol=1e-4)
      PY

Gotchas

• Dynamic shapes need dynamic_axes, or the exported graph fixes the batch size.
• Set the model to eval() first so dropout and batchnorm behave deterministically.

Related guides