Module 5 — Graph optimization

An ONNX file straight out of torch.onnx.export or tf2onnx is a literal transcription of the training graph. Every layer sits in its own node; every intermediate tensor is materialised; every constant tensor computed at build time is stored as a runtime Constant node evaluated at every inference. That is exactly what we want for verification — the exported graph should reflect the trained model — and exactly what we do not want for performance. This module explains what ONNX Runtime does when it loads that graph, how to control the transformations, and how to measure the resulting speedup on the ResNet18 and the text encoder.

The interesting part is that these optimizations are free: the artifact is unchanged, the numerics stay within tolerance, and the code that calls the runtime does not change. Only the internal graph, invisible to the caller, becomes leaner.

What ONNX Runtime does at session creation

Loading a .onnx file into an InferenceSession triggers a pipeline of graph transformations before the first inference runs. Three families of transformations dominate the gain.

Constant folding. Nodes whose inputs are all constants — a shape computation, a normalisation weight, a scalar tensor — get evaluated once at load time and replaced by their result. The rest of the graph stops depending on them. On a ResNet, this typically removes a dozen Cast and Reshape nodes that were computing shape information at every request.

Operator fusion. Consecutive operators that map to a single efficient kernel get merged. A canonical example: a Conv followed by a BatchNormalization followed by a Relu becomes a fused Conv+BN+Relu node. The intermediate tensors between them are never materialised, memory bandwidth drops, and cache locality improves. On the ResNet, module 5 typically fuses more than 20 such patterns.

Redundancy elimination. Identity nodes get removed. Two consecutive Transpose nodes that cancel each other become nothing. A Cast from float32 to float32 disappears. These sound trivial but accumulate: a naive export can leave dozens of them, each costing a memory copy.

Together, these transformations often shrink the ResNet18 graph from around 60 nodes to below 30. That is not by itself the speedup, but it is a strong indicator.

Optimization levels

ONNX Runtime exposes four levels through SessionOptions.graph_optimization_level, in increasing aggressiveness:

Level	Meaning
ORT_DISABLE_ALL	No optimization. The graph runs exactly as exported. Useful only for debugging.
ORT_ENABLE_BASIC	Constant folding, common subexpression elimination, layout-independent rewrites. Always safe.
ORT_ENABLE_EXTENDED	Level 1 plus operator fusion. The default in most releases.
ORT_ENABLE_ALL	Level 2 plus layout-specific rewrites: NCHW to NCHWc on CPU, transformer-specific fusion patterns. The recommended level for production.

The default (ORT_ENABLE_ALL on recent releases) is the right choice for production. ORT_DISABLE_ALL exists for one specific case: when the numerical parity check of module 4 fails and you want to bisect the cause. Turning optimizations off will tell you whether the bug lives in the exporter or in a runtime rewrite.

import onnxruntime as ort

so = ort.SessionOptions()
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
session = ort.InferenceSession("resnet.onnx", so, providers=["CPUExecutionProvider"])

Saving the optimized graph

The pipeline runs on every session creation. On a small model this is milliseconds; on a large transformer it is seconds, and a service that restarts frequently pays that cost every time. optimized_model_filepath lets ONNX Runtime write the post-optimization graph to disk, then load that graph directly on future startups:

so = ort.SessionOptions()
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
so.optimized_model_filepath = "resnet.optimized.onnx"

_ = ort.InferenceSession("resnet.onnx", so, providers=["CPUExecutionProvider"])

The _ is intentional: the session gets created, the optimized file gets written, and we throw the session away. Subsequent serving processes load resnet.optimized.onnx directly, skipping the pipeline. Open the two files in Netron side by side: the optimized graph is visibly denser, with fused blocks and constants inlined into the operators that need them.

Optimizer output is not portable across execution providers

The optimized graph baked out with the CPU provider active contains CPU-specific layout choices. Loading it with the CUDA provider works but discards some of those choices; TensorRT ignores it entirely. Ship the original .onnx as the source of truth and let each runtime produce its own optimized cache.

Measuring the gain

The right way to measure the speedup is with the protocol from module 8: warm-up runs, then several hundred timed inferences, then percentiles. In outline, on a laptop CPU:

import time, numpy as np
import onnxruntime as ort

def benchmark(session, x, n=200, warmup=20):
    for _ in range(warmup):
        session.run(None, {"input": x})
    ts = []
    for _ in range(n):
        t0 = time.perf_counter()
        session.run(None, {"input": x})
        ts.append(time.perf_counter() - t0)
    return np.median(ts) * 1000

x = np.random.rand(1, 3, 224, 224).astype(np.float32)

so_off = ort.SessionOptions()
so_off.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
so_on = ort.SessionOptions()
so_on.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

off = benchmark(ort.InferenceSession("resnet.onnx", so_off, providers=["CPUExecutionProvider"]), x)
on  = benchmark(ort.InferenceSession("resnet.onnx", so_on,  providers=["CPUExecutionProvider"]), x)

print(f"no optim: {off:.2f} ms | full optim: {on:.2f} ms | speedup: {off/on:.2f}x")

On the ResNet18, expect a factor of 1.3 to 1.8x on CPU, mostly from Conv+BN+Relu fusion. On the text encoder, the gain concentrates on the attention pattern — MatMul, Softmax, MatMul fused into Attention — and typically reaches 2 to 3x when the runtime recognises the pattern (transformer-specific fusion requires level ORT_ENABLE_ALL).

Numerical parity survives all four levels for a correctly exported model. Re-run the check of module 4 with each graph_optimization_level; the max absolute difference should stay under 1e-4. If it grows with the level, one of the fusion rewrites is misbehaving on this model — a rare occurrence, but the reason each fusion pass can be individually disabled via the optimizer.disable_... session flags.

Constant folding, explicitly

onnxsim is a companion project that runs a stronger constant-folding pass before the runtime sees the file. It is useful in two cases: when the runtime version at deployment time is fixed and does not include a rewrite that the optimizer would want, and when the file will be shipped to a mobile runtime with limited optimization capability. The command line is a single call:

pip install onnxsim
python -m onnxsim resnet.onnx resnet.folded.onnx

The output typically has 5 to 15 fewer nodes than the input. Run the parity check afterwards, same as always — a simplified graph is a new artifact and deserves its own verification.

In summary

• ONNX Runtime rewrites the graph on every session creation: constant folding, operator fusion (Conv+BN+Relu and friends), and redundancy elimination are the three families that dominate the gain.
• GraphOptimizationLevel has four steps; ORT_ENABLE_ALL is the production default, ORT_DISABLE_ALL exists to bisect a numerical bug.
• Save the optimized graph with optimized_model_filepath to skip the pipeline on subsequent startups; keep the original as the shipped artifact because the optimizer output is provider-specific.
• Measure the gain with the module 8 protocol: expect 1.3–2x on the ResNet from operator fusion, 2–3x on transformers when attention fusion kicks in, always followed by a re-run of module 4's parity check.

Next module: quantization, where we trade precision for size and speed, and where the numerical parity check earns its keep.