Running Model Inference over Large Rasters

A trained model that works beautifully on 256-pixel chips has to be applied to scenes of 10,980 × 10,980 pixels or to mosaics a hundred times larger. Doing that correctly is a raster engineering problem with three requirements that pull against each other: the output must be pixel-aligned with the input, no tile boundary may be visible in the result, and peak memory must stay bounded regardless of scene size. This topic belongs to Raster Machine Learning & Model Inference and picks up where building training datasets from satellite imagery leaves off.

The core idea is a halo: read more than you write. Each tile is read with a margin of context around it, the model predicts the whole padded tile, and only the centre is written. Every output pixel therefore had a full neighbourhood available when it was predicted, and the tile edges disappear.

---

Prerequisites

pip install "rasterio>=1.3.0" "numpy>=1.23" "onnxruntime>=1.16" "tqdm>=4.66"

Conceptually, you need a model artifact that declares its band order and scaling, as set out on the section overview, and a source raster that is internally tiled — a striped GeoTIFF forces a full-width read for every window and destroys the performance of this whole pattern. Converting the input is covered in converting a GeoTIFF to a COG with rio-cogeo.

---

Step-by-Step Workflow

Step 1 — Plan the read and write windows

Every tile has two windows: the write window, which tiles the scene exactly with no gaps and no overlap, and the read window, which is the write window grown by the halo and clipped to the scene. The write windows partition the output; the read windows overlap each other.

from dataclasses import dataclass

import rasterio
from rasterio.windows import Window


@dataclass(frozen=True)
class Tile:
    read: Window
    write: Window
    crop: tuple[int, int, int, int]   # top, bottom, left, right trim in read pixels


def plan_tiles(width: int, height: int, size: int = 256, halo: int = 64):
    """Write windows tile the scene exactly; read windows add a clipped halo."""
    for row in range(0, height, size):
        for col in range(0, width, size):
            w = min(size, width - col)
            h = min(size, height - row)
            write = Window(col, row, w, h)
            r0, c0 = max(0, row - halo), max(0, col - halo)
            r1 = min(height, row + h + halo)
            c1 = min(width, col + w + halo)
            read = Window(c0, r0, c1 - c0, r1 - r0)
            yield Tile(read, write, (row - r0, r1 - (row + h), col - c0, c1 - (col + w)))

Storing the crop amounts on the tile removes the arithmetic from the hot loop, and — more importantly — makes the edge cases explicit. A tile at the top-left of the scene has no halo above or to the left, so its crop is (0, 64, 0, 64) rather than (64, 64, 64, 64).

Step 2 — Derive the output profile from the input

The output must sit on the source grid. The reliable way to guarantee that is never to construct a transform at all: copy the source profile and override only what changes.

import rasterio


def open_prediction_writer(src: rasterio.DatasetReader, path: str,
                           *, n_classes: int = 1):
    profile = src.profile | {
        "count": n_classes,
        "dtype": "uint8",
        "nodata": 0,
        "compress": "deflate",
        "predictor": 2,        # horizontal differencing suits class rasters
        "tiled": True,
        "blockxsize": 512,
        "blockysize": 512,
        "BIGTIFF": "IF_SAFER",
    }
    profile.pop("photometric", None)   # inherited RGB hints break a 1-band output
    return rasterio.open(path, "w", **profile)

Dropping photometric matters when the source is a three-band true-colour product: GDAL will otherwise refuse a single-band write with an inherited photometric=RGB. The compression choices are the same ones weighed in choosing COG compression: ZSTD vs DEFLATE, with the caveat that class rasters compress extraordinarily well under predictor=2.

Step 3 — Batch reads so the model is never starved

A model call on a batch of sixteen tiles costs barely more than a call on one, because the framework overhead is amortised. The read, however, scales linearly. Batching the reads and overlapping them with compute is what turns a half-utilised device into a saturated one.

import numpy as np


def batched_tiles(src, tiles, batch_size: int = 16, scale: float = 10_000.0):
    """Yield (arrays, tiles) batches of padded tiles read from the source."""
    buf_arr, buf_tile = [], []
    for tile in tiles:
        arr = src.read(window=tile.read, boundless=True,
                       fill_value=src.nodata or 0).astype("float32") / scale
        buf_arr.append(arr)
        buf_tile.append(tile)
        if len(buf_arr) == batch_size:
            yield buf_arr, buf_tile
            buf_arr, buf_tile = [], []
    if buf_arr:
        yield buf_arr, buf_tile

Note that the arrays in a batch are not all the same shape: edge tiles have a smaller halo. Either pad them to the nominal read size before stacking, or group tiles by shape. Padding is simpler and the cost is negligible because edge tiles are a small fraction of any real scene.

Step 4 — Predict, crop, write

The crop is where the halo is discarded. Getting the sign of the trim wrong here produces an output that is shifted by exactly the halo width — a bug that looks like a georeferencing problem and is not.

import numpy as np


def run_inference(session, src, dst, tiles, batch_size: int = 16) -> None:
    input_name = session.get_inputs()[0].name
    for arrays, batch in batched_tiles(src, tiles, batch_size):
        padded, shapes = pad_to_common(arrays)          # (n, bands, H, W)
        logits = session.run(None, {input_name: padded})[0]
        classes = logits.argmax(axis=1).astype("uint8") + 1
        for pred, tile, (h, w) in zip(classes, batch, shapes):
            top, bottom, left, right = tile.crop
            core = pred[:h, :w][top:h - bottom, left:w - right]
            dst.write(core, 1, window=tile.write)

The + 1 after argmax keeps class 0 reserved for “not predicted”, matching the convention used throughout this section. The nodata mask is applied on the same core array before writing when the source has fill pixels.

Step 5 — Finish the file

A prediction raster is a data product, and a data product needs overviews. For a class raster those overviews must use nearest-neighbour or mode resampling — averaging class codes produces values that correspond to no class at all.

import rasterio
from rasterio.enums import Resampling

with rasterio.open("prediction.tif", "r+") as dst:
    dst.build_overviews([2, 4, 8, 16, 32], Resampling.mode)
    dst.update_tags(ns="rio_overview", resampling="mode")
    dst.update_tags(
        model_name="landcover_v3",
        model_contract_sha="9f1c2a…",
        source_scene="S2A_36NYF_20260614",
    )

Writing the model identity into the file’s tags is what makes the output traceable months later; the mechanics are the same as in reading and writing GDAL tags and band descriptions.

---

Scaling the Loop Across Processes and Machines

One process streaming tiles will saturate a laptop’s disk and a single GPU. Beyond that, the loop has to spread out, and the shape it takes depends entirely on which model class you are running.

Pixel-based models parallelise by window with no coordination at all. Each worker opens its own dataset handle — GDAL dataset objects are emphatically not safe to share across processes — takes a slice of the window list, and writes into its own output file. Merging the per-worker outputs afterwards is a plain mosaic of non-overlapping pieces, which is the cheapest case of merging tiles with rasterio merge.

from concurrent.futures import ProcessPoolExecutor

import rasterio


def worker(args) -> str:
    src_path, dst_path, tiles = args
    with rasterio.open(src_path) as src:          # open INSIDE the worker
        with open_prediction_writer(src, dst_path) as dst:
            run_inference(load_session(), src, dst, tiles)
    return dst_path


def run_parallel(src_path: str, tiles: list, n_workers: int = 8) -> list[str]:
    chunks = [tiles[i::n_workers] for i in range(n_workers)]
    jobs = [(src_path, f"pred_part{i}.tif", c) for i, c in enumerate(chunks)]
    with ProcessPoolExecutor(max_workers=n_workers) as pool:
        return list(pool.map(worker, jobs))

Striping the tile list with tiles[i::n_workers] rather than handing each worker a contiguous block is deliberate for a different reason than load balance: it spreads each worker’s reads across the whole file, so no two workers are fighting over the same byte ranges of a remote object at the same moment.

Convolutional models on a shared GPU want the opposite arrangement — one process owning the device, several reader threads feeding it. Threads work here where processes would not, because the read releases the GIL inside GDAL and the model call releases it inside the framework.

When a single machine can no longer finish the backlog overnight, the unit of distribution becomes the scene rather than the tile. One scene per task, each task self-contained, is the granularity that survives retries and partial failures, and it maps directly onto the job patterns in distributed processing on Coiled and AWS Batch. Resist the temptation to distribute individual tiles across machines: the scheduling overhead per tile swamps the tens of milliseconds of work, and the output assembly becomes a distributed write.

---

Parameter Reference

---

Verification & Testing

The cheapest verification is a grid assertion: the prediction must have exactly the source’s shape and transform.

import rasterio

with rasterio.open("s2_stack_36NYF.tif") as src, \
     rasterio.open("prediction.tif") as pred:
    assert pred.transform == src.transform, "output grid drifted"
    assert (pred.width, pred.height) == (src.width, src.height)
    assert pred.crs == src.crs

The second check is a seam test. Read a strip that crosses several tile boundaries and look at the class-change frequency as a function of column. If changes spike at exactly every 256th column, the halo is not being applied — either it is zero, or the crop arithmetic is discarding the wrong side.

import numpy as np
import rasterio

with rasterio.open("prediction.tif") as pred:
    strip = pred.read(1, window=((5000, 5001), (0, pred.width)))[0]
changes = np.flatnonzero(np.diff(strip) != 0)
at_seams = (changes % 256 == 0).mean()
print(f"{at_seams:.1%} of class changes fall exactly on a tile boundary")

A healthy run puts well under 1% of changes on tile boundaries. A broken run puts 20% or more there. This test has the great virtue of needing no reference data — it detects the defect from the output alone. The visual equivalent, and the fix when it fails, is in stitching prediction tiles without seams.

Finally, check the class histogram against expectations before publishing anything. A land-cover product where 60% of pixels are class 0 means the nodata mask swallowed the scene, usually because the scale factor was applied twice and every feature landed outside the model’s training range.

---

Troubleshooting

The output is shifted by exactly the halo width

The crop trims from the wrong side, or boundless padding was included in the core slice. Assert on a tile in the middle of the scene that core.shape equals (tile.write.height, tile.write.width) before writing; a mismatch shows up immediately.

rasterio.errors.WindowError: Bounds and transform are inconsistent

A read window with negative offsets reached src.read without boundless=True. Clip the read window to the dataset with Window(0, 0, src.width, src.height).intersection(read) or pass boundless=True and a fill value.

Inference is fast on the first hundred tiles then collapses

The output file is being written out of block order and GDAL is rewriting blocks. Iterate row-major over the write grid, keep the write block size a divisor of the tile size, and never write the same block twice.

onnxruntime raises about a dynamic dimension

The exported model has a fixed batch dimension and the last batch is short. Export with a dynamic batch axis, or pad the final batch up to batch_size and discard the extra predictions.

Memory climbs steadily through the run

Predictions are being accumulated instead of written, or the progress bar is holding references to arrays. Write inside the loop, and never build a list of per-tile outputs — the whole point of the pattern is that nothing larger than a batch is resident.

---

Frequently Asked Questions

Q: How much overlap does a segmentation model need? At least half the receptive field of the network, rounded up to a multiple of the downsampling factor. For a four-level U-Net that is typically 32 to 64 pixels per side. Too little overlap leaves seams; too much multiplies the compute for no gain.

Q: Why is my GPU only twenty per cent utilised during inference? Because reading is serialised with compute. A single loop that reads a tile, predicts it, then writes it leaves the GPU idle for the whole read. Put the reader on a thread feeding a bounded queue so the next batch is decoded while the current one is on the device.

Q: Can I write predictions straight into a Cloud-Optimized GeoTIFF? Write a tiled GeoTIFF block by block, then add overviews and re-order the header. Writing tiled output first and converting once at the end is far faster than trying to maintain COG layout while random blocks arrive.

Q: Does a pixel-based model need a halo at all? No. A model that sees one pixel at a time has a receptive field of one pixel, so set halo=0 and read exactly the write window. The halo only exists to give convolutions their context.

---

Related

• Tiled Inference with Overlapping Windows — the halo pattern in full, with edge handling.
• Stitching Prediction Tiles without Seams — blending probabilities where cropping alone is not enough.
• Batching GPU Inference from a Rasterio Window Reader — the reader thread, the queue, and the utilisation numbers.
• Running ONNX Models on Raster Tiles — exporting once and running anywhere without a training stack.
• Optimizing Rasterio Window Reads for Memory Efficiency — the read discipline this whole loop rests on.