Common ONNX export failures
Why this matters
Exporting to ONNX is how you move models from research to production inference engines (TensorRT, ONNX Runtime, CoreML), but the export process masks errors that only surface at inference time in production, making debugging extremely costly.
Explanation
What it is: ONNX export converts a PyTorch model to an interchange format that other frameworks can run. The export process uses torch.onnx.export(), which traces your model by running it with dummy inputs. This tracing approach creates a static computation graph: but many PyTorch patterns don't map to static graphs. How it works mechanically: When you call torch.onnx.export(), PyTorch executes your forward pass with a special tracing mode. Control flow (if statements based on tensor values), missing operator implementations in ONNX, and in-place operations that modify module state all break the static graph assumption. The export completes without error, but the exported model either runs incorrectly or fails at runtime in the target framework. When to use this knowledge: Before any production ONNX deployment, you must validate that your model's architecture avoids dynamic control flow, uses only standard operations, and doesn't rely on stateful patterns.
Analogy
Exporting to ONNX is like taking a snapshot of your model's computation at one frozen moment. If your model's behavior changes based on the actual data values (not just shapes), the snapshot captured one path, but different inputs will try to follow a different path that wasn't recorded.
Code
import torch
import torch.nn as nn
import torch.onnx
import tempfile
import os
class DynamicControlFlowModel(nn.Module):
def forward(self, x):
if x.sum() > 0:
return x * 2
else:
return x * 3
class MissingOpModel(nn.Module):
def forward(self, x):
return torch.nonzero(x)
class StatefulModel(nn.Module):
def __init__(self):
super().__init__()
self.counter = 0
def forward(self, x):
self.counter += 1
return x * self.counter
def attempt_export(model, model_name):
dummy_input = torch.randn(2, 3)
with tempfile.TemporaryDirectory() as tmpdir:
onnx_path = os.path.join(tmpdir, f"{model_name}.onnx")
try:
torch.onnx.export(
model,
dummy_input,
onnx_path,
input_names=["input"],
output_names=["output"],
opset_version=14
)
print(f"{model_name}: Export succeeded (but may have issues at runtime)")
except Exception as e:
print(f"{model_name}: Export failed with {type(e).__name__}: {str(e)[:80]}")
print("=== Testing Dynamic Control Flow ===")
model1 = DynamicControlFlowModel()
attempt_export(model1, "DynamicControlFlow")
print(" Issue: if/else based on tensor value → not traceable")
print()
print("=== Testing Missing Op ===")
model2 = MissingOpModel()
attempt_export(model2, "MissingOp")
print(" Issue: torch.nonzero not in ONNX opset 14")
print()
print("=== Testing Stateful Model ===")
model3 = StatefulModel()
attempt_export(model3, "Stateful")
print(" Issue: self.counter mutation not captured in static graph")
print()
print("=== Safe Alternative: Refactored Model ===")
class SafeModel(nn.Module):
def forward(self, x):
mask = (x.sum() > 0).float()
return x * (2.0 * mask + 3.0 * (1 - mask))
model4 = SafeModel()
attempt_export(model4, "Safe")
print(" Fix: Replace if/else with tensor operations")=== Testing Dynamic Control Flow === DynamicControlFlow: Export succeeded (but may have issues at runtime) Issue: if/else based on tensor value → not traceable === Testing Missing Op === MissingOp: Export failed with RuntimeError: ONNX export failed: Couldn't export operator aten::nonzero Issue: torch.nonzero not in ONNX opset 14 === Testing Stateful Model === Stateful: Export succeeded (but may have issues at runtime) Issue: self.counter mutation not captured in static graph === Safe Alternative: Refactored Model === Safe: Export succeeded (but may have issues at runtime) Fix: Replace if/else with tensor operations
What just happened?
The code demonstrates three failure modes: (1) Dynamic control flow exports without error but produces wrong results because ONNX records only the traced path; (2) Missing operators fail explicitly at export time because torch.nonzero doesn't exist in ONNX opset 14; (3) Stateful mutations export without error but the exported model loses the counter state because ONNX doesn't support in-place module attribute updates. The safe refactor replaces if/else with element-wise tensor operations that map directly to ONNX ops.
Common gotcha
The most dangerous failure mode is silent success: the export completes without error, you move the model to production, and it produces wrong results because the static graph only recorded one branch of control flow. Always validate the exported model's outputs match PyTorch's on a diverse set of test cases, not just the dummy input.
Error recovery
RuntimeError: Couldn't export operator aten::XValueError: Tracing failed / input must be TensorModel output shape/values don't match PyTorchExperienced dev note
The reason ONNX export seems to succeed when it shouldn't: PyTorch's tracing mode executes your forward() with a special proxy tensor that records operations, not actual values. When you have an if statement that checks a tensor's value, the tracer arbitrarily picks one branch (usually True), records it, and exports a graph that always takes that path. Your model works fine in PyTorch because Python evaluates the real condition each time. In ONNX, you've hardcoded one path forever. Use torch.onnx.export(..., verbose=True) to see what was actually traced: the verbose output shows which ops got recorded and reveals when only one branch of a conditional is present.
Check your understanding
You export a model that classifies inputs as 'positive' or 'negative' using `if x.mean() > threshold:` in the forward method. The exported ONNX model gives correct results on positive samples but wrong results on negative samples. Why, and what's the minimal fix?
Show answer hint
A correct answer identifies that only the positive branch was traced (the if-condition was true for the dummy input), so negative samples follow the wrong code path in ONNX. The minimal fix is to refactor the if statement into a single tensor-based computation that computes both branches and selects one based on a mask, like using torch.where().