Exporting and Verifying Stateful PyTorch Sequence Models with ONNX Runtime

Hello! Last lesson used future-period and instrument-level results to choose a GRU checkpoint for its generalization, rather than its training loss. The next question is operational: can you run that same checkpoint outside PyTorch without changing its predictions?

Today you will export a stateful GRU step to ONNX and compare it with ONNX Runtime. The key is to test not just one prediction, but a sequence of predictions and the hidden state passed between them. That gives you a dependable starting point for the deployment checks in the next lesson.

Make the state an explicit part of the model

ONNX stores a computation and its parameters; ONNX Runtime executes that computation. A hidden state kept privately in a Python object is not a reliable interface between runtime calls. For deployment, make the state a tensor input and a tensor output:

The middle dimension of x is sequence length because the example below uses batch_first=True. Batch size and sequence length are both fixed at one in this first export. That matches a step-by-step inference loop and avoids introducing dynamic-shape problems before correctness is established. Stateful here means the caller passes h_out back as h_in for the next step; ONNX Runtime does not remember it automatically.

For a short introduction to what the exported file and runtime each do, watch the opening of ONNX Explained with Example, by Daniel Krei:

A graph viewer may show a GRU surrounded by shape and tensor-manipulation operations. The supplied graph fragment is an example, not the graph you should expect your exporter to produce:

In particular, the pictured GRU has weight dimensions involving $384=3\times 128$: a GRU maintains three gates per hidden unit. You do not need to reproduce those nodes by hand. Your responsibility is to define and verify the graph’s external tensor contract.

Export the checkpoint, not a newly initialized model

PyTorch’s Export a PyTorch model to ONNX tutorial shows the current torch.onnx.export(..., dynamo=True) workflow, saving the resulting program, and checking the file:

Here is a compact inference interface for a one-direction GRU with a two-output head. Adapt the class and checkpoint loading to the model you trained: F, H, L, the head, and all preprocessing must match that checkpoint. The numerical values below are illustrative, not competition feature specifications.

import torch
from torch import nn

class GRUStep(nn.Module):
    def __init__(self, feature_count, hidden_size, num_layers):
        super().__init__()
        self.gru = nn.GRU(
            input_size=feature_count,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
        )
        self.head = nn.Linear(hidden_size, 2)

    def forward(self, x, h_in):
        # x: [1, 1, F]; h_in: [L, 1, H]
        sequence_out, h_out = self.gru(x, h_in)
        prediction = self.head(sequence_out[:, 0, :])
        return prediction, h_out

F, H, L = 48, 128, 1  # Replace with the trained model's configuration.
model = GRUStep(F, H, L)
state_dict = torch.load(
    "selected_state_dict.pt", map_location="cpu", weights_only=True
)
model.load_state_dict(state_dict, strict=True)
model.eval()

x0 = torch.zeros(1, 1, F, dtype=torch.float32)
h0 = torch.zeros(L, 1, H, dtype=torch.float32)

with torch.no_grad():
    y0, h1 = model(x0, h0)
assert y0.shape == (1, 2)
assert h1.shape == (L, 1, H)

onnx_program = torch.onnx.export(
    model,
    (x0, h0),
    input_names=["x", "h_in"],
    output_names=["prediction", "h_out"],
    dynamo=True,
    verify=True,
)
onnx_program.save("gru_step.onnx", external_data=False)

import onnx
onnx.checker.check_model(onnx.load("gru_step.onnx"))

strict=True helps catch a mismatch between your model definition and saved weights. Calling eval() makes inference behavior explicit; it matters if your trained model contains dropout. The exported file includes the loaded parameters. Keep your feature preparation identical to validation: this interface expects already prepared float32 features and does not export preprocessing.

The exporter documentation distinguishes its automatic verify=True check from saving the ONNX program. It also documents dynamic shapes, which you deliberately do not need for this fixed one-step interface:

verify=True is useful, but it does not prove that state is handled correctly over a stream. If the recommended exporter fails on your particular GRU and PyTorch version, inspect its error or request an export report before changing the model. A practical compatibility option is the older exporter:

# Use only if the recommended export path fails for your GRU/version.
torch.onnx.export(
    model, (x0, h0), "gru_step.onnx",
    input_names=["x", "h_in"],
    output_names=["prediction", "h_out"],
    dynamo=False,
)

Check and test the resulting file in exactly the same way below. Do not silently accept a different export path because it created a file; successful serialization is not numerical validation.

Verify two evolving computations

The PyTorch tutorial demonstrates how ONNX Runtime takes NumPy arrays keyed by input name and how to compare its outputs with PyTorch:

Use two separately evolving states during the test. PyTorch must receive its own previous h_out, while ONNX Runtime receives its previous h_out. Feeding PyTorch’s state into both on every step would test individual calls, but could hide an error that accumulates when ONNX Runtime runs independently.

import numpy as np
import onnxruntime as ort
import torch

session = ort.InferenceSession(
    "gru_step.onnx", providers=["CPUExecutionProvider"]
)
assert [v.name for v in session.get_inputs()] == ["x", "h_in"]
assert [v.name for v in session.get_outputs()] == [
    "prediction", "h_out"
]

rng = np.random.default_rng(7)
streams = [
    rng.standard_normal((13, F)).astype(np.float32),
    rng.standard_normal((5, F)).astype(np.float32),
]
initial_states = [
    np.zeros((L, 1, H), dtype=np.float32),
    rng.standard_normal((L, 1, H)).astype(np.float32),
]

max_prediction_error = 0.0
max_state_error = 0.0

for stream, initial_h in zip(streams, initial_states):
    # A new stream starts from its specified state, never the previous
    # stream's final state. In ordinary independent streams, use zeros.
    h_torch = torch.from_numpy(initial_h.copy())
    h_ort = initial_h.copy()

    for row in stream:
        x_np = row.reshape(1, 1, F)

        with torch.no_grad():
            y_torch, h_torch = model(
                torch.from_numpy(x_np), h_torch
            )

        y_ort, h_ort_next = session.run(
            None, {"x": x_np, "h_in": h_ort}
        )

        assert y_ort.shape == (1, 2)
        assert h_ort_next.shape == (L, 1, H)
        assert np.isfinite(y_ort).all()
        assert np.isfinite(h_ort_next).all()

        max_prediction_error = max(
            max_prediction_error,
            np.max(np.abs(y_torch.numpy() - y_ort)),
        )
        max_state_error = max(
            max_state_error,
            np.max(np.abs(h_torch.numpy() - h_ort_next)),
        )

        torch.testing.assert_close(
            y_torch, torch.from_numpy(y_ort),
            atol=1e-4, rtol=1e-4,
        )
        torch.testing.assert_close(
            h_torch, torch.from_numpy(h_ort_next),
            atol=1e-4, rtol=1e-4,
        )

        h_ort = h_ort_next

print("Maximum prediction error:", max_prediction_error)
print("Maximum hidden-state error:", max_state_error)

The nonzero initial state in the second synthetic stream tests whether the exported graph really uses h_in; it is not the normal reset policy. For independent validation sequences, initialize each state to zeros and process every input step in order, including steps without a scored target. Repeat the comparison on representative preprocessed, held-out streams, including long ones. Log the step and sequence where any comparison fails. A short random stream is a smoke test; it does not establish agreement over the inputs you intend to serve.

atol protects comparisons near zero, where relative error is uninformative; rtol allows error proportional to the output’s magnitude. The values above are initial float32 test tolerances, not a license to ignore large discrepancies. Compare both outputs and inspect maximum errors. A hidden-state difference can precede a prediction difference, especially after many steps.

What a passing export does—and does not—establish

The checks have distinct jobs. onnx.checker.check_model checks that the file is structurally valid. Creating an ONNX Runtime session checks that this runtime can load it. Numerical comparisons test whether this saved file, under the tested input and state conditions, agrees with the selected PyTorch checkpoint.

If comparisons fail, first inspect mundane interface errors: checkpoint mismatch, feature order or scaling, float32 conversion, input and output order, state shape, and whether both implementations reset at the same true sequence boundaries. If the first step agrees but later steps diverge, inspect h_out closely. If both implementations agree but the validation score changes, investigate the surrounding preprocessing, masks, stream ordering, and scoring code rather than assuming the ONNX graph is at fault.

The single-step graph is intentionally a correctness baseline. It preserves the two predictions per step and explicit state passing without requiring a long unrolled graph. It does not yet establish that a complete submission meets its memory, file-size, CPU, or time limits.

Takeaways

A deployable stateful GRU has a clear four-tensor contract: current features and incoming state produce two predictions and outgoing state. Export the selected, evaluation-mode checkpoint, check the saved graph, then run PyTorch and ONNX Runtime independently through multi-step streams while comparing both predictions and hidden states. Reset each implementation at the same real sequence boundaries.

Next, you will profile and package this verified inference path against the submission constraints.