Fixing Opset and Dynamic Shape Errors During ONNX Export

Resolve ONNX export failures for geospatial models: unsupported opset operators, dynamic tile sizes, batch axes, ZipMap output surprises and numerical parity checks.

Three settings decide whether an export is usable: the opset version (low enough for the deployed runtime, high enough for the operators), the dynamic axes (batch, and height/width for convolutional models), and — for scikit-learn classifiers — turning off ZipMap so the output is a tensor rather than a list of dictionaries. Get those right and the remaining work is a numerical parity check.

This page is the troubleshooting guide. For why to export at all, and how to serve the result, see ONNX Export for Geospatial Model Inference.

The opset must sit inside the window your runtime supports A number line of opset versions is shown. The operators used by the model require at least opset thirteen. The runtime deployed in the inference container supports up to opset seventeen. Only versions thirteen to seventeen are usable, and exporting at opset twenty produces a file the container cannot load. Export target = highest opset the deployed runtime accepts 91113 172022 operators unsupported below 13 usable window runtime in the container is too old min opset from ops max opset from runtime Exporting at the newest available version is the most common cause of a “works on my laptop, fails in the container” deployment.

Why This Fails in Geospatial ML Pipelines

The opset mismatch is a deployment-time failure, which is the worst kind. Export with the newest torch on a workstation and the file targets opset 20; the inference image pins onnxruntime==1.14 for GDAL compatibility and supports up to 17. The export succeeded, the file is valid, and the container raises InvalidGraph on load — at 4 a.m., in the tile task, after the sensor has already waited three hours for the scene.

The second failure is a shape baked in by constant folding. dynamic_axes marks an axis as free, but if the model computes something from x.shape[2] in Python rather than with a tensor operation, the tracer records the concrete value. The export runs, the parity check on the same 256-pixel patch passes, and the first 512-pixel tile in production fails with a shape mismatch. Testing two different shapes is the only way to catch it.

Third, ZipMap. skl2onnx wraps classifier probabilities in a sequence of maps, so sess.run(...) returns a list of dictionaries rather than an array. In a per-pixel raster loop that is both a type surprise and a serious slowdown — millions of dictionary allocations per tile. Disabling it changes the output to a plain tensor and typically makes inference several times faster.

Core Principles

  • Pin the opset from the runtime, not from the exporter.
  • Mark every axis that varies — batch, height, width — and verify with two shapes.
  • Disable ZipMap for scikit-learn classifiers.
  • Check parity numerically, with a tolerance you chose deliberately.
  • Test-load the artifact with the deployment runtime in CI.
  • Record the opset and runtime version in the model card.

Production-Ready Code

from __future__ import annotations

import logging
import numpy as np
import onnx
import onnxruntime as ort

logger = logging.getLogger(__name__)


def runtime_max_opset() -> int:
    """Highest opset the INSTALLED onnxruntime supports."""
    domains = ort.get_available_providers() and onnx.defs.onnx_opset_version()
    return int(domains)


def export_torch(model, sample: "torch.Tensor", path: str, opset: int) -> None:
    """Export a convolutional model with dynamic batch, height and width."""
    import torch

    model.eval()
    torch.onnx.export(
        model, sample, path,
        input_names=["input"], output_names=["logits"],
        dynamic_axes={"input": {0: "batch", 2: "height", 3: "width"},
                      "logits": {0: "batch", 2: "height", 3: "width"}},
        opset_version=opset, do_constant_folding=True,
    )
    onnx.checker.check_model(onnx.load(path))
    logger.info("exported %s at opset %d", path, opset)


def export_sklearn(pipeline, n_features: int, path: str, opset: int) -> None:
    """Export a scikit-learn classifier WITHOUT ZipMap, with a dynamic row axis."""
    from skl2onnx import to_onnx
    from skl2onnx.common.data_types import FloatTensorType

    initial = [("input", FloatTensorType([None, n_features]))]
    model = to_onnx(pipeline, initial_types=initial, target_opset=opset,
                    options={id(pipeline): {"zipmap": False}})
    with open(path, "wb") as fh:
        fh.write(model.SerializeToString())
    onnx.checker.check_model(onnx.load(path))
    logger.info("exported %s at opset %d (zipmap disabled)", path, opset)


def assert_dynamic(path: str, axes: tuple[int, ...] = (0, 2, 3)) -> None:
    """Fail if an axis that should be free was baked to a constant."""
    graph = onnx.load(path).graph
    dims = graph.input[0].type.tensor_type.shape.dim
    for axis in axes:
        d = dims[axis]
        if d.HasField("dim_value"):
            raise ValueError(
                f"axis {axis} is fixed at {d.dim_value}; it should be dynamic. "
                "A Python-level use of x.shape was traced as a constant.")


def check_parity(path: str, reference_fn, sample: np.ndarray,
                 atol: float = 1e-4) -> float:
    """Max absolute difference between the ONNX graph and the source model."""
    sess = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
    got = sess.run(None, {sess.get_inputs()[0].name: sample.astype("float32")})[0]
    want = np.asarray(reference_fn(sample))
    if got.shape != want.shape:
        raise ValueError(f"shape mismatch: onnx {got.shape} vs source {want.shape}")
    diff = float(np.max(np.abs(got - want)))
    if diff > atol:
        raise ValueError(f"numerical parity failed: max |diff| = {diff:.3e} > {atol}")
    logger.info("parity ok, max |diff| = %.3e", diff)
    return diff


def check_two_shapes(path: str, base: np.ndarray) -> None:
    """The real dynamic-shape test: two different sizes, one session."""
    sess = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
    name = sess.get_inputs()[0].name
    small = base[:, :, :128, :128].astype("float32")
    large = base[:, :, :512, :512].astype("float32")
    for arr in (small, large):
        out = sess.run(None, {name: arr})[0]
        assert out.shape[-2:] == arr.shape[-2:], f"shape not preserved for {arr.shape}"

Step-by-Step Walkthrough

Step 1 — Pick the opset from the deployment image, not the laptop

docker run --rm geoml-inference:2026.08 python -c \
  "import onnxruntime, onnx; print(onnxruntime.__version__, onnx.defs.onnx_opset_version())"
# 1.18.0 21

Record that number in configuration. Export targets min(runtime_max, chosen) and never the exporter’s default.

Step 2 — Export, then immediately assert the axes

import torch

TARGET_OPSET = 17
sample = torch.zeros(1, 6, 256, 256)
export_torch(model, sample, "artifacts/unet_v7.onnx", opset=TARGET_OPSET)
assert_dynamic("artifacts/unet_v7.onnx", axes=(0, 2, 3))

assert_dynamic is the cheap check that catches constant folding before it reaches production.

Step 3 — Run parity and the two-shape test

import numpy as np

x = np.random.rand(2, 6, 256, 256).astype("float32")
check_parity("artifacts/unet_v7.onnx",
             reference_fn=lambda a: model(torch.from_numpy(a)).detach().numpy(),
             sample=x, atol=1e-4)
check_two_shapes("artifacts/unet_v7.onnx", x)

A parity failure at 1e-4 on a convolutional model usually means the model was exported in training mode — check model.eval() — because batch-norm running statistics differ between modes.

Step 4 — For scikit-learn, kill ZipMap and confirm the output type

export_sklearn(pipeline, n_features=12, path="artifacts/landcover_rf.onnx",
               opset=TARGET_OPSET)

sess = ort.InferenceSession("artifacts/landcover_rf.onnx")
out = sess.run(None, {sess.get_inputs()[0].name:
                      np.random.rand(5, 12).astype("float32")})
print([type(o) for o in out], out[1].shape)
assert isinstance(out[1], np.ndarray), "ZipMap is still enabled"

Step 5 — Test-load in the deployment image in CI

- name: verify the artifact loads in the inference runtime
  run: |
    docker run --rm -v "$PWD/artifacts:/a" geoml-inference:2026.08 \
      python -c "import onnxruntime as ort; \
                 s=ort.InferenceSession('/a/unet_v7.onnx'); \
                 print([i.shape for i in s.get_inputs()])"

This one CI step removes the entire class of “valid file, wrong runtime” incidents.

How a dynamic axis quietly becomes a constant Model code that reads the input height as a Python integer and uses it in an arithmetic expression is traced during export, so the concrete value two hundred and fifty-six is written into the graph. The exported model then accepts only that size, even though the axis was declared dynamic. dynamic_axes declares intent; the tracer decides reality traced as a constant h = x.shape[2] pad = (32 − h % 32) % 32 Python int → baked into the graph stays symbolic use nn.functional.pad with a tensor-derived size, or avoid shape arithmetic entirely input shape [1, 6, 256, 256] any other tile size raises at run time input shape [batch, 6, height, width] assert_dynamic passes; two shapes both run Inspect graph.input[0] after every export — it takes one line and tells you the truth.

Verification

import numpy as np
import onnx
import onnxruntime as ort
import pytest


def test_opset_is_within_the_runtime_window():
    model = onnx.load("artifacts/unet_v7.onnx")
    exported = max(o.version for o in model.opset_import if o.domain in ("", "ai.onnx"))
    assert exported <= onnx.defs.onnx_opset_version()


def test_axes_are_dynamic():
    assert_dynamic("artifacts/unet_v7.onnx", axes=(0, 2, 3))


def test_two_tile_sizes_both_run():
    x = np.random.rand(1, 6, 512, 512).astype("float32")
    check_two_shapes("artifacts/unet_v7.onnx", x)


def test_parity_within_tolerance():
    x = np.random.rand(2, 6, 256, 256).astype("float32")
    diff = check_parity("artifacts/unet_v7.onnx",
                        reference_fn=torch_reference, sample=x, atol=1e-4)
    assert diff < 1e-4


def test_sklearn_output_is_a_tensor():
    sess = ort.InferenceSession("artifacts/landcover_rf.onnx")
    out = sess.run(None, {sess.get_inputs()[0].name:
                          np.random.rand(4, 12).astype("float32")})
    assert all(isinstance(o, np.ndarray) for o in out), "ZipMap still enabled"


def test_nodata_rows_do_not_change_valid_predictions():
    """A NaN row must not contaminate its batch neighbours."""
    sess = ort.InferenceSession("artifacts/landcover_rf.onnx")
    name = sess.get_inputs()[0].name
    x = np.random.rand(4, 12).astype("float32")
    alone = sess.run(None, {name: x[:1]})[0]
    x_with_nan = x.copy()
    x_with_nan[3] = np.nan
    together = sess.run(None, {name: x_with_nan})[0][:1]
    assert np.allclose(alone, together, atol=1e-5)

The last test is worth keeping for tree models in particular: some export paths handle NaN differently from the source library, and a raster batch always contains nodata rows.

Six checks, in this order, before the artifact ships A left to right sequence of six checks: choose the opset from the deployed runtime, export with dynamic axes declared, assert that the axes really are dynamic, check numerical parity against the source model, run two different input shapes, and finally load the artifact inside the deployment image. Each step catches a failure the previous one cannot 1. opset from the image 2. export dynamic_axes 3. assert axes are free 4. parity max |diff| 5. two shapes 128 and 512 6. load in the image Steps 3–6 take under a minute combined and belong in CI, not in a runbook. Skipping step 6 is how a valid file reaches production and fails to load.

FAQ

Which opset version should I target?

The highest one the deployed runtime accepts, provided it covers your operators. Read it from the inference image, pin it in configuration, and never take the exporter’s default.

How do I export a model that must accept any tile size?

Mark height and width — and batch, if you batch tiles — in dynamic_axes, then verify with two different shapes. Constant folding can bake a shape in even when the axis is declared dynamic.

Why does my exported classifier return a list of dictionaries?

That is skl2onnx’s ZipMap. Disable it with options={id(pipeline): {"zipmap": False}} so the output is a plain probability tensor — both simpler and considerably faster in a raster loop.

Part of: ONNX Export for Geospatial Model Inference Part of: Geospatial MLOps and Model Deployment