Running ONNX Models on Raster Tiles

Export once with dynamic axes, then run the model from a container that has no training framework in it at all:

import numpy as np
import onnxruntime as ort

sess = ort.InferenceSession("landcover_v3.onnx",
                            providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
name = sess.get_inputs()[0].name

logits = sess.run(None, {name: tile.astype("float32")[None]})[0]   # (1, C, H, W)
classes = logits.argmax(axis=1).astype("uint8") + 1

The reason to bother is operational rather than numerical: the inference image shrinks by an order of magnitude and the graph stops changing under you. This page extends running model inference over large rasters in Raster Machine Learning & Model Inference.

Why a Frozen Graph Suits Raster Pipelines

A raster inference job runs on whatever machine the scheduler gives it, often hundreds of times, often months after the model was trained. Every dependency in that image is a chance for the prediction to drift.

What each inference image has to carry A container with the full training framework and CUDA libraries is roughly six gigabytes and pulls in hundreds of packages. The same pipeline on ONNX Runtime with GDAL is around four hundred megabytes with a handful of packages, which starts faster on a batch worker and has far fewer moving parts to pin. Inference image weight training stack ~6.0 GB · 300+ packages · 90 s cold start ONNX + GDAL ~400 MB · 12 packages · 8 s cold start On a batch queue the cold start is paid per task, not per run. A thousand scene-level tasks turn 80 seconds of difference into a day of compute.

The exported graph also freezes the operations, so an upgrade to the training library on a cluster image cannot alter a prediction. For anything that will be re-run to produce a consistent time series, that stability is worth more than the convenience of keeping the original objects.

Environment & Setup

Package Version pin Used for
onnxruntime-gpu >=1.16 The inference session; use onnxruntime for CPU-only
rasterio >=1.3.0 Window reads and writes around the session
numpy >=1.23 Array handling at the session boundary
torch >=2.1 Export only; not needed in the inference image 
pip install "onnxruntime-gpu>=1.16" "rasterio>=1.3.0" "numpy>=1.23"

Complete Working Example

Export first, with the contract written into the file:

import json

import onnx
import torch


def export_with_contract(model, path: str, *, bands: list[str], scale: float,
                         classes: dict[int, str], chip: int = 512) -> None:
    model.eval()
    dummy = torch.zeros(1, len(bands), chip, chip)
    torch.onnx.export(
        model, dummy, path,
        input_names=["reflectance"], output_names=["logits"],
        dynamic_axes={"reflectance": {0: "batch", 2: "height", 3: "width"},
                      "logits": {0: "batch", 2: "height", 3: "width"}},
        opset_version=17,
    )
    graph = onnx.load(path)
    for key, value in {
        "band_order": json.dumps(bands),
        "reflectance_scale": str(scale),
        "classes": json.dumps(classes),
    }.items():
        entry = graph.metadata_props.add()
        entry.key, entry.value = key, value
    onnx.save(graph, path)

Then run it, asserting the contract against the raster before a single tile is read:

import json

import numpy as np
import onnxruntime as ort
import rasterio


def open_session(path: str):
    sess = ort.InferenceSession(
        path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
    meta = sess.get_modelmeta().custom_metadata_map
    contract = {
        "bands": json.loads(meta["band_order"]),
        "scale": float(meta["reflectance_scale"]),
        "classes": {int(k): v for k, v in json.loads(meta["classes"]).items()},
    }
    return sess, contract


def check_source(src: rasterio.DatasetReader, contract: dict) -> None:
    names = [src.descriptions[i] for i in range(src.count)]
    if names != contract["bands"]:
        raise ValueError(f"band mismatch: file {names} vs model {contract['bands']}")


def predict_tiles(sess, contract, arrays: list[np.ndarray]) -> np.ndarray:
    """Pad to a common shape, run one batch, return uint8 classes."""
    h = max(a.shape[1] for a in arrays)
    w = max(a.shape[2] for a in arrays)
    batch = np.zeros((len(arrays), len(contract["bands"]), h, w), dtype="float32")
    for i, a in enumerate(arrays):
        batch[i, :, :a.shape[1], :a.shape[2]] = a / contract["scale"]

    name = sess.get_inputs()[0].name
    logits = sess.run(None, {name: batch})[0]
    return logits.argmax(axis=1).astype("uint8") + 1

Padding to a common shape rather than relying on dynamic axes for every distinct edge size is the pragmatic choice. Dynamic axes make the export tolerant, but a runtime that sees twelve different input shapes re-plans its memory twelve times; padding to one shape keeps a single plan warm.

Providers, Precision and Determinism

Choosing an execution provider If a CUDA device is present and the model is convolutional, the CUDA provider is the right choice. Without a device, or for a shallow model where transfer dominates, the CPU provider with thread tuning is faster in practice. TensorRT is worth the extra build time only for a model that will run over many thousands of scenes unchanged. Provider selection in practice Is a CUDA device visible? no — CPU provider set intra_op threads to cores minus the readers yes — CUDA provider the default answer for any convolutional model TensorRT only if the same graph runs over thousands of scenes Provider order in the list is a preference, not a guarantee — always log which one the session actually chose.

The provider list is a preference order, and a session will silently fall back to CPU if the CUDA libraries are missing from the image. That failure mode looks like “inference got fifty times slower last Tuesday”, so log sess.get_providers() at startup and fail loudly if the expected provider is absent.

Precision is the other lever. float16 export roughly halves the model time on tensor-core hardware, at the cost of small differences in the logits. For a classification output those differences matter only where the top two classes are within a fraction of a per cent of each other — real, but confined to pixels that were ambiguous anyway. Where the product is a probability rather than a class, keep float32.

Determinism deserves a decision rather than an assumption. The same ONNX file on the same provider produces identical output, which is what makes reruns reproducible; the same file across providers does not, because reduction orders differ. If a time series is being assembled over months, pin the provider as well as the model version, and record both in the output tags alongside the other provenance described in attaching model metadata to STAC items.

Verifying the Export before It Goes to Production

An export is a translation, and translations need a test. The test that matters is not whether the file loads but whether it produces the same decisions as the model it came from.

Three checks between export and production First the session must load and report the expected provider. Second the logits from one validation tile must match the source framework to within a small tolerance. Third the class decisions across a validation set must agree at better than ninety-nine point nine per cent, which is the check that catches a genuinely broken export. Gate the export on all three 1 · session loads metadata map present, expected provider active catches packaging faults 2 · logits agree max absolute difference under 1e-3 on one tile catches opset drift 3 · classes agree over 99.9% of pixels on a held-out validation set the one that must pass Check two can fail harmlessly on a float16 export; check three cannot fail harmlessly at all. Run all three in CI on every export, not once by hand at the end of the project. 
import numpy as np


def verify_export(torch_model, sess, tile: np.ndarray, atol: float = 1e-3) -> None:
    import torch

    with torch.inference_mode():
        ref = torch_model(torch.from_numpy(tile[None])).numpy()
    got = sess.run(None, {sess.get_inputs()[0].name: tile[None]})[0]

    assert ref.shape == got.shape, (ref.shape, got.shape)
    print("max |diff| :", float(np.abs(ref - got).max()))
    agree = (ref.argmax(1) == got.argmax(1)).mean()
    assert agree > 0.999, f"class agreement only {agree:.4%}"

Run this in continuous integration on every export, against a small committed tile, and the whole class of “the exported model is subtly different” incidents disappears. The tile can be synthetic as long as it exercises the full value range; what matters is that the comparison happens automatically rather than in someone’s notebook.

Common Errors

INVALID_ARGUMENT: Got invalid dimensions for input

The tile shape does not match a static axis. Re-export with the spatial dimensions in dynamic_axes, or pad every tile to the exported size.

The session runs on CPU despite a GPU being present

The CUDA libraries in the image do not match the onnxruntime-gpu build. Check ort.get_available_providers() and pin both the runtime and the CUDA base image, using the reproducibility patterns in pinning GDAL and PROJ versions reproducibly.

Predictions differ slightly from the training framework

Expected, and usually irrelevant: opset conversion changes operation order and the results differ in the last few significant digits. Verify by comparing class agreement on a validation tile rather than by comparing floats; anything below 99.9% agreement indicates a real export problem.

Frequently Asked Questions

Q: Why export to ONNX rather than ship the training framework? Size and stability. An inference image with ONNX Runtime and GDAL is a few hundred megabytes against several gigabytes with a full training stack, and the exported graph is frozen, so a framework upgrade on the cluster cannot silently change your predictions.

Q: How do I keep the band contract with the model file? ONNX files carry a metadata map. Write band order, reflectance scale and class codes into it at export time and read them back at session creation, so the model and its contract cannot be separated.

Q: Do dynamic spatial axes cost performance? A little, because the runtime cannot pre-plan every allocation. In practice the loss is small next to the flexibility of handling edge tiles, and it disappears entirely if you pad every tile to one nominal size.