ONNX Export for Geospatial Model Inference

Training and inference rarely share the same home. A land-cover model is trained in a data-science environment stuffed with scikit-learn, XGBoost, PyTorch, CUDA, and a matching Python version — then it must run inside a lean container that reads Cloud-Optimized GeoTIFFs, streams tiles, and writes a classified raster back to object storage. Shipping the entire training stack into that inference image is slow to build, fragile to upgrade, and painful to reproduce two years later when a security patch forces a dependency bump. The Open Neural Network Exchange (ONNX) format breaks that coupling: you export a fitted model to a single self-describing .onnx graph, then serve it with ONNX Runtime, a compact engine with no dependency on the framework that produced the model.

This is part of Geospatial MLOps and Model Deployment, which covers the path from a trained model to a monitored production service. This guide walks through exporting the three model families you actually deploy on raster data — a scikit-learn or XGBoost land-cover classifier and a PyTorch convolutional network — declaring the correct input tensor shapes, wiring dynamic axes for variable tile sizes, and proving numerical parity before you trust the artifact in production.

Problem Framing — why ONNX for geospatial inference

Geospatial inference environments are dominated by GDAL. Reading a Cloud-Optimized GeoTIFF, warping a tile, and writing a georeferenced output all route through the GDAL/PROJ system libraries that rasterio wraps. That stack is already heavy and version-sensitive. Layering the full training toolchain on top — a specific scikit-learn build, XGBoost with its native library, PyTorch with a CUDA runtime — turns a container image into a multi-gigabyte liability whose builds break every time one pin shifts. ONNX cuts the training frameworks out of the serving image entirely.

Three properties make the format the right fit for raster work:

• Environment decoupling. The .onnx artifact carries the full computation graph and weights. ONNX Runtime loads it without scikit-learn, XGBoost, or PyTorch installed, so your inference image is a slim GDAL layer plus onnxruntime. This is the foundation of a reproducible containerized inference pipeline.
• Speed. ONNX Runtime applies graph optimizations — constant folding, operator fusion, and layout tuning — and dispatches to hardware-specific execution providers. For tree ensembles and convolutional networks it routinely beats the native Python prediction path, which matters when you push millions of pixels through a model.
• Portability. One artifact runs on a CPU laptop, a CUDA server, or an edge box with an integrated GPU, selecting the provider at load time. The model you validate locally is byte-for-byte the model you deploy.

The models you export here are the same ones covered elsewhere on this site: the tree ensembles from gradient boosting for raster data and the neural architectures related to graph neural networks for spatial data. Export is the bridge between training them and running them in production, and it sits alongside model drift detection and spatial dataset versioning in a complete deployment story.

Prerequisites & Environment Setup

Export is where version drift bites hardest. The opset your exporter emits must be understood by the onnxruntime that loads the file, and skl2onnx must match the scikit-learn that trained the model. Pin everything.

# Pinned toolchain for ONNX export and inference
scikit-learn==1.5.1
xgboost==2.1.1
skl2onnx==1.17.0
onnxmltools==1.12.0          # XGBoost -> ONNX bridge
onnx==1.16.2
onnxconverter-common==1.14.0 # shared opset/type helpers
onnxruntime==1.19.2
torch==2.4.1
numpy==1.26.4
rasterio==1.3.10             # raster IO at inference time

Install the export environment:

pip install "scikit-learn==1.5.1" "xgboost==2.1.1" "skl2onnx==1.17.0" \
            "onnxmltools==1.12.0" "onnx==1.16.2" "onnxconverter-common==1.14.0" \
            "onnxruntime==1.19.2" "torch==2.4.1" "numpy==1.26.4" "rasterio==1.3.10"

Confirm the opset ceiling your runtime supports before you export anything. Emitting an opset newer than the runtime understands is the single most common cause of a model that converts cleanly but fails to load in production:

import onnxruntime as ort
import onnx

print("onnxruntime:", ort.__version__)
# The default ONNX domain opset supported by this build
domain_versions = ort.get_available_providers()
print("providers:", domain_versions)

TARGET_OPSET = 17   # safe modern baseline for skl2onnx and torch.onnx.export

Treat TARGET_OPSET as a shared constant across every exporter in this guide. A single mismatched opset between the scikit-learn and PyTorch exports produces two artifacts that need different runtime builds — exactly the fragmentation ONNX is meant to remove.

Step-by-Step Implementation

Step 1 — Export a scikit-learn or XGBoost land-cover classifier

A land-cover classifier consumes a fixed number of spectral bands per pixel — say six Sentinel-2 bands — and returns a class label. The whole export hinges on one decision: the FloatTensorType input shape. It must be [None, n_bands], where None is a dynamic batch (the number of pixels you score in a call) and n_bands is fixed to the feature count the model was trained on.

import numpy as np
from sklearn.ensemble import RandomForestClassifier
from skl2onnx import to_onnx
from skl2onnx.common.data_types import FloatTensorType

N_BANDS = 6   # Sentinel-2 bands used as features, must match training

def export_sklearn_classifier(model, n_bands: int, path: str, opset: int = 17) -> None:
    """Export a fitted scikit-learn classifier to ONNX.

    Args:
        model: A fitted scikit-learn estimator (e.g. RandomForestClassifier).
        n_bands: Number of input features (spectral bands) per pixel.
        path: Output .onnx file path.
        opset: Target ONNX opset version.
    """
    if not hasattr(model, "classes_"):
        raise ValueError("Model must be fitted before export (no classes_ found)")
    # None -> dynamic batch (pixels per call); n_bands -> fixed feature count
    initial_type = [("float_input", FloatTensorType([None, n_bands]))]
    onnx_model = to_onnx(
        model,
        initial_types=initial_type,
        target_opset=opset,
        options={id(model): {"zipmap": False}},   # return a plain probability array
    )
    with open(path, "wb") as f:
        f.write(onnx_model.SerializeToString())

# Train a small classifier on stacked band values, one row per pixel
X = np.random.rand(5000, N_BANDS).astype(np.float32)
y = np.random.randint(0, 5, size=5000)   # 5 land-cover classes
clf = RandomForestClassifier(n_estimators=200, max_depth=12, random_state=42).fit(X, y)

export_sklearn_classifier(clf, N_BANDS, "landcover_rf.onnx")

The zipmap=False option is not cosmetic. By default skl2onnx wraps classifier probabilities in a ZipMap operator that emits a list of dictionaries — awkward to consume and slow to vectorize over a raster. Disabling it returns a clean 2-D probability array you can reshape straight back to a raster grid.

XGBoost does not go through skl2onnx; it uses onnxmltools, which understands the booster’s native tree structure:

import numpy as np
from xgboost import XGBClassifier
from onnxmltools import convert_xgboost
from onnxmltools.convert.common.data_types import FloatTensorType as MLToolsFloat

def export_xgboost_classifier(model, n_bands: int, path: str, opset: int = 17) -> None:
    """Export a fitted XGBoost classifier to ONNX via onnxmltools."""
    initial_type = [("float_input", MLToolsFloat([None, n_bands]))]
    onnx_model = convert_xgboost(model, initial_types=initial_type, target_opset=opset)
    with open(path, "wb") as f:
        f.write(onnx_model.SerializeToString())

X = np.random.rand(5000, N_BANDS).astype(np.float32)
y = np.random.randint(0, 5, size=5000)
xgb = XGBClassifier(n_estimators=300, max_depth=8, tree_method="hist").fit(X, y)

export_xgboost_classifier(xgb, N_BANDS, "landcover_xgb.onnx")

Both artifacts now expect a float32 array of shape [pixels, 6] and return class probabilities. If you fitted the classifier on a scaled feature space, bake the same scaling into your preprocessing before inference — the ONNX graph only holds the estimator, not the feature scaling you applied upstream.

Step 2 — Export a PyTorch CNN with dynamic axes

A convolutional land-cover model works on tiles, not pixel rows. Its input is a 4-D tensor [batch, channels, height, width]. Production tiling almost never uses one fixed size — edge tiles are ragged, and you may reblock between runs — so the export must accept variable spatial dimensions. That is exactly what dynamic_axes provides.

import torch
import torch.nn as nn

class SmallSegNet(nn.Module):
    """A compact fully-convolutional network for per-pixel land-cover classification."""
    def __init__(self, in_channels: int = 6, n_classes: int = 5):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_channels, 32, 3, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(),
            nn.Conv2d(64, n_classes, 1),   # 1x1 conv -> class logits per pixel
        )

    def forward(self, x):
        return self.net(x)

def export_cnn(model, in_channels: int, path: str, opset: int = 17) -> None:
    """Export a fully-convolutional CNN to ONNX with dynamic batch and tile size."""
    model.eval()
    # Trace with a representative tile; concrete values do not constrain dynamic axes
    dummy = torch.randn(1, in_channels, 256, 256)
    torch.onnx.export(
        model,
        dummy,
        path,
        input_names=["tile"],
        output_names=["logits"],
        opset_version=opset,
        dynamic_axes={
            "tile":   {0: "batch", 2: "height", 3: "width"},
            "logits": {0: "batch", 2: "height", 3: "width"},
        },
        do_constant_folding=True,
    )

model = SmallSegNet(in_channels=6, n_classes=5)
export_cnn(model, in_channels=6, path="landcover_cnn.onnx")

The dynamic_axes dictionary marks axes 0, 2, and 3 as free on both input and output. Without it, the exporter bakes the traced 256×256 shape into the graph and every other tile size raises a shape-mismatch error at inference. Note that a fully convolutional network is what makes variable spatial dimensions valid; if the architecture contains a flatten into a dense layer, only the batch axis can be dynamic, and you must tile to a fixed size.

Step 3 — Validate parity between original and ONNX outputs

Never trust an export you have not checked. Run the original model and the ONNX graph on identical input and compare. For the scikit-learn classifier:

import numpy as np
import onnxruntime as ort

def check_sklearn_parity(model, onnx_path: str, n_bands: int, n: int = 512) -> None:
    """Assert the ONNX probabilities match the scikit-learn model."""
    X = np.random.rand(n, n_bands).astype(np.float32)
    ref = model.predict_proba(X).astype(np.float32)

    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
    input_name = sess.get_inputs()[0].name
    onnx_out = sess.run(None, {input_name: X})[1]   # index 1 = probabilities
    onnx_out = np.asarray(onnx_out, dtype=np.float32)

    assert np.allclose(ref, onnx_out, atol=1e-4), \
        f"Parity failed: max abs diff = {np.abs(ref - onnx_out).max():.2e}"
    print(f"sklearn parity OK — max abs diff {np.abs(ref - onnx_out).max():.2e}")

check_sklearn_parity(clf, "landcover_rf.onnx", N_BANDS)

For the CNN, feed the same tile through PyTorch (in eval mode, under no_grad) and through ONNX Runtime, exercising a non-traced tile size to prove the dynamic axes work:

import numpy as np
import torch
import onnxruntime as ort

def check_cnn_parity(model, onnx_path: str, in_channels: int) -> None:
    """Assert the ONNX CNN output matches PyTorch on a non-traced tile size."""
    model.eval()
    # Deliberately not 256x256 — proves dynamic height/width axes work
    tile = torch.randn(2, in_channels, 320, 192)
    with torch.no_grad():
        ref = model(tile).numpy()

    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
    onnx_out = sess.run(None, {"tile": tile.numpy().astype(np.float32)})[0]

    assert np.allclose(ref, onnx_out, atol=1e-4), \
        f"CNN parity failed: max abs diff = {np.abs(ref - onnx_out).max():.2e}"
    print(f"CNN parity OK on {tuple(tile.shape)} — max abs diff {np.abs(ref - onnx_out).max():.2e}")

check_cnn_parity(model, "landcover_cnn.onnx", in_channels=6)

If either assertion fails, do not deploy. Parity failure means the exported graph computes something different from your validated model, and the classified raster will be silently wrong.

Verification & Testing

Wire parity into the export step so a broken artifact never reaches a registry. A single helper that loads the model, checks the declared input signature, and asserts numerical agreement gives you a gate you can run in CI.

import numpy as np
import onnx
import onnxruntime as ort

def verify_onnx_artifact(onnx_path: str, ref_fn, sample: np.ndarray,
                         output_index: int = 0, atol: float = 1e-4) -> None:
    """Full artifact gate: structural check + numerical parity.

    Args:
        onnx_path: Path to the exported .onnx file.
        ref_fn: Callable returning the reference output for `sample`.
        sample: A representative float32 input array.
        output_index: Which session output to compare.
        atol: Absolute tolerance for np.allclose.
    """
    # 1. Structural validity — raises if the graph is malformed
    onnx.checker.check_model(onnx.load(onnx_path))

    # 2. Load and confirm the input dtype is float32, not float64
    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
    declared = sess.get_inputs()[0].type
    assert "float" in declared, f"Unexpected input type {declared}"
    assert sample.dtype == np.float32, "Input must be float32 to match FloatTensorType"

    # 3. Numerical parity against the source model
    ref = np.asarray(ref_fn(sample), dtype=np.float32)
    got = np.asarray(sess.run(None, {sess.get_inputs()[0].name: sample})[output_index],
                     dtype=np.float32)
    max_diff = np.abs(ref - got).max()
    assert np.allclose(ref, got, atol=atol), f"Parity failed: max abs diff {max_diff:.2e}"
    print(f"Artifact verified — max abs diff {max_diff:.2e}")

sample = np.random.rand(256, N_BANDS).astype(np.float32)
verify_onnx_artifact("landcover_rf.onnx", clf.predict_proba, sample, output_index=1)

The gate checks three failure modes at once: a structurally invalid graph, a silent float64 input (which raises at runtime, not export), and numerical drift from the source model. Passing all three is the contract that lets a downstream service load the artifact without re-validating.

Troubleshooting & Common Errors

Unsupported model IR version or missing operator at load

Cause: The exporter emitted an opset newer than the deployed onnxruntime understands. The model converts fine on your workstation but fails when the production container loads it.

Fix: Set target_opset (skl2onnx / onnxmltools) or opset_version (torch) to a value the runtime supports, and pin both packages. Opset 17 loads on onnxruntime>=1.15. Never let the exporter pick a default opset implicitly.

RuntimeException: Unexpected input data type. Actual: (tensor(double))

Cause: You declared a FloatTensorType (float32) input but passed a float64 array. NumPy defaults raster reads and arithmetic to float64, so this is the most frequent geospatial export bug.

Fix: Cast every input with .astype(np.float32) before session.run. Rasterio returns the raster’s native dtype, so cast immediately after reading a tile:

tile = src.read().astype(np.float32)   # cast before inference, always

Shape error on a tile that is not the traced size

Cause: A CNN exported without dynamic_axes has the traced height and width frozen into the graph, so any other tile size is rejected.

Fix: Re-export with dynamic_axes marking axes 2 and 3 as dynamic, as in Step 2. Confirm with a parity check on a deliberately different tile shape. If the network is not fully convolutional, you must tile to the fixed trained size instead.

ValueError: The maximum opset needed by this model is only N

Cause: A component of the model — often a preprocessing step folded into an sklearn Pipeline — has no converter for the requested opset, or the estimator type is not registered with skl2onnx.

Fix: Export the estimator without the unsupported preprocessing step and reproduce that preprocessing in NumPy before inference. Alternatively, register a custom converter. Keep the ONNX graph to the model itself and handle scaling and band math outside it — see raster band math and index calculation for reproducing index features at inference time.

Probabilities returned as a list of dictionaries

Cause: The default skl2onnx ZipMap operator wraps classifier probabilities in per-sample dictionaries, which cannot be reshaped into a raster.

Fix: Pass options={id(model): {"zipmap": False}} at export, as in Step 1, to get a plain 2-D probability array.

Unsupported operator during CNN export

Cause: A layer or op in the PyTorch model has no ONNX equivalent at the chosen opset — custom autograd functions and some exotic interpolation modes are common offenders.

Fix: Raise the opset if a newer one adds the operator, replace the layer with an ONNX-supported equivalent (for example, standard nn.Upsample modes), or split the op out of the graph and run it in NumPy. Export a minimal sub-model to isolate which layer fails.

Performance Optimisation

Choose the right execution provider

ONNX Runtime dispatches to execution providers in the order you list them, falling back to the next when an operator is unsupported. Put the accelerator first and CPU last as a guaranteed fallback:

import onnxruntime as ort

def make_session(onnx_path: str, use_gpu: bool = False) -> ort.InferenceSession:
    """Create an inference session with a sensible provider fallback chain."""
    if use_gpu:
        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
    else:
        providers = ["CPUExecutionProvider"]
    opts = ort.SessionOptions()
    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    opts.intra_op_num_threads = 0   # 0 = let ORT pick based on cores
    return ort.InferenceSession(onnx_path, sess_options=opts, providers=providers)

Setting graph_optimization_level to ORT_ENABLE_ALL triggers constant folding and operator fusion at load, which is free throughput for both tree and CNN models.

Stream tiles instead of loading whole rasters

A continental raster does not fit in memory, and you do not need it to. Read windows with rasterio, run each window through the session, and write the result back tile by tile. This keeps memory bounded regardless of raster size:

import numpy as np
import rasterio
from rasterio.windows import Window

def infer_raster_tiled(onnx_path: str, in_path: str, out_path: str,
                       tile: int = 512, n_bands: int = 6) -> None:
    """Run tiled ONNX inference over a raster, writing a classified output."""
    sess = make_session(onnx_path)
    input_name = sess.get_inputs()[0].name

    with rasterio.open(in_path) as src:
        profile = src.profile.copy()
        profile.update(count=1, dtype="uint8")
        with rasterio.open(out_path, "w", **profile) as dst:
            for row in range(0, src.height, tile):
                for col in range(0, src.width, tile):
                    win = Window(col, row,
                                 min(tile, src.width - col),
                                 min(tile, src.height - row))
                    block = src.read(window=win).astype(np.float32)  # (bands, h, w)
                    bands, h, w = block.shape
                    # reshape to (pixels, bands) for a tree/pixel classifier
                    X = block.reshape(bands, h * w).T
                    proba = sess.run(None, {input_name: X})[1]
                    labels = np.asarray(proba).argmax(axis=1).astype("uint8")
                    dst.write(labels.reshape(h, w), 1, window=win)

For a CNN, skip the reshape and feed the [1, bands, h, w] tile directly, thanks to the dynamic axes from Step 2. Mapping predictions back to the correct grid cells is the same discipline covered in mapping XGBoost predictions back to a raster grid.

Batch pixels and reuse the session

Session creation is expensive; inference is cheap. Build the session once and reuse it across every tile rather than reopening the file per window. For tree ensembles, larger pixel batches amortize the per-call overhead — feeding one full 512×512 tile as 262144×6 rows is far faster than looping row by row.

FAQ

Does ONNX preserve exact predictions from the original model?

Tree ensembles convert with near-exact parity, typically within 1e-5 absolute error, because the decision logic maps directly onto ONNX tree operators. Neural networks match to float32 precision. The small residual differences come from float32 rounding and fused operations, so validate with np.allclose using an atol around 1e-4 rather than expecting bit-identical output.

Which opset version should I target for geospatial models?

Target the highest opset your deployed onnxruntime supports. Opset 17 is a safe modern baseline for both skl2onnx and torch.onnx.export. Pin the opset explicitly at export time and match it to the runtime version so no operator is ever missing at load. Emitting a default opset implicitly is how the same model passes on a workstation and fails in the container.

Can one ONNX file handle raster tiles of different sizes?

Yes, for fully convolutional CNNs. Declare height and width as dynamic axes in torch.onnx.export, and the exported graph accepts any tile shape the network supports. Fixed-shape exports only accept the exact dimensions traced at export time and raise a shape error on any other tile. Tree and pixel classifiers are naturally size-agnostic because they see a [pixels, bands] matrix with a dynamic batch.

Do I still need GDAL and rasterio in the inference container?

You need rasterio or GDAL to read and write the georeferenced raster, but not the training stack. ONNX Runtime replaces scikit-learn, XGBoost, and PyTorch at inference time, so the container ships a slim raster IO layer plus onnxruntime instead of the heavy training dependencies — the core saving that ONNX export delivers.

Related

• Model Drift Detection for Geospatial Inference — monitor the deployed ONNX model for covariate shift in input bands
• Containerizing Geospatial Inference Pipelines — package the ONNX artifact and a slim GDAL layer into a reproducible image
• Spatial Dataset Versioning with DVC and lakeFS — version the training data behind each exported artifact
• Serving an ONNX Land-Cover Model with ONNX Runtime — a complete inference service around the exported model
• Gradient Boosting for Raster Data — train the XGBoost classifier this guide exports

Part of: Geospatial MLOps and Model Deployment