Batching GPU Inference from a Rasterio Window Reader

To stop a GPU idling while rasterio reads, move the reads onto threads that fill a bounded queue and consume batches from it:

import queue, threading

q: queue.Queue = queue.Queue(maxsize=4)          # bounded: readers block when ahead

def reader():
    for tile in tiles:
        q.put((src.read(window=tile.read, boundless=True, fill_value=0), tile))
    q.put(None)

threading.Thread(target=reader, daemon=True).start()
while (item := q.get()) is not None:
    arrays, tile = item
    write(model(arrays))                          # device stays busy

GDAL releases the GIL during decompression, so a reader thread genuinely overlaps with the model call rather than merely interleaving with it. This page extends running model inference over large rasters in Raster Machine Learning & Model Inference.

Where the Idle Time Comes From

A naive loop alternates three phases that cannot overlap: read and decompress, run the model, crop and write. For a 640-pixel six-band tile from a compressed COG the read is tens of milliseconds, the model call on a modern GPU a few, and the write a handful. The device is therefore idle for most of the run, and buying a faster card changes nothing.

Device utilisation against reader threads With no separate reader the device is busy about a quarter of the time. One reader thread lifts it to roughly half, two to three quarters, and four to just under ninety per cent. Beyond four threads the curve is flat because the storage rather than the reader is now the limit. Reader threads versus GPU utilisation 100% 0% 25% 50% 75% 89% 91% 91% 0 1 2 4 8 16 Local NVMe, 640 px six-band tiles. Object storage needs more threads to reach the same plateau.

Environment & Setup

Package Version pin Used for
rasterio >=1.3.0 Thread-safe windowed reads from per-thread dataset handles
torch >=2.1 Pinned memory, non-blocking transfers, inference mode
numpy >=1.23 Batch stacking 
pip install "rasterio>=1.3.0" "torch>=2.1" "numpy>=1.23"

Complete Working Example

import queue
import threading

import numpy as np
import rasterio
import torch


def reader_worker(src_path: str, tiles: list, out_q: queue.Queue,
                  scale: float = 10_000.0) -> None:
    """One dataset handle per thread — handles must not be shared."""
    with rasterio.open(src_path) as src:
        fill = src.nodata if src.nodata is not None else 0
        for tile in tiles:
            arr = src.read(window=tile.read, boundless=True, fill_value=fill)
            out_q.put((arr.astype("float32") / scale, tile))
    out_q.put(None)


def batched(out_q: queue.Queue, n_readers: int, batch_size: int):
    """Yield batches, ending when every reader has signalled completion."""
    buf, done = [], 0
    while done < n_readers:
        item = out_q.get()
        if item is None:
            done += 1
            continue
        buf.append(item)
        if len(buf) == batch_size:
            yield buf
            buf = []
    if buf:
        yield buf


@torch.inference_mode()
def run(model, src_path: str, dst_path: str, tiles: list, *,
        batch_size: int = 16, n_readers: int = 4, device: str = "cuda") -> None:
    q: queue.Queue = queue.Queue(maxsize=batch_size * 3)
    chunks = [tiles[i::n_readers] for i in range(n_readers)]
    for chunk in chunks:
        threading.Thread(target=reader_worker, args=(src_path, chunk, q),
                         daemon=True).start()

    model = model.to(device).eval()
    with rasterio.open(src_path) as src:
        profile = src.profile | {"count": 1, "dtype": "uint8", "nodata": 0,
                                 "compress": "deflate", "tiled": True}
        profile.pop("photometric", None)
        with rasterio.open(dst_path, "w", **profile) as dst:
            for batch in batched(q, n_readers, batch_size):
                arrays = np.stack([a for a, _ in batch])
                # pin_memory + non_blocking overlaps the host-to-device copy
                x = torch.from_numpy(arrays).pin_memory().to(device, non_blocking=True)
                pred = model(x).argmax(dim=1).to("uint8").cpu().numpy() + 1

                for out, (_, tile) in zip(pred, batch):
                    t, b, l, r = tile.crop
                    h, w = out.shape
                    dst.write(out[t:h - b, l:w - r], 1, window=tile.write)
Readers, queue, device, writer Four reader threads each hold their own dataset handle and push decoded tiles into a bounded queue. The main thread drains the queue, stacks tiles into a batch, transfers it to the device and runs the model, then crops each prediction and writes it to the destination file. The queue bound propagates back-pressure to the readers. One process, four readers, one device reader 1 — own handle reader 2 — own handle reader 3 — own handle reader 4 — own handle bounded queue maxsize caps memory and blocks readers batch and transfer stack 16 tiles pin, copy, run argmax on device crop and write trim per-tile halo write block-aligned nothing accumulates Back-pressure flows right to left: a slow write eventually stalls the readers, which is the desired behaviour.

The structure has one more benefit that only shows up on long runs: it fails cleanly. Because nothing is buffered beyond the queue bound and every prediction is written as soon as it is cropped, killing the process loses at most one batch of work, and the partially written output is a valid raster covering the blocks completed so far.

The bounded queue is what makes this safe. An unbounded queue lets fast readers run ahead of a slow model until the process is holding thousands of tiles and the machine swaps; maxsize makes the readers block instead, which is exactly the back-pressure you want.

One dataset handle per thread is non-negotiable. A rasterio.DatasetReader is not thread-safe, and sharing one produces corrupted reads that look like random noise in a handful of tiles — a bug that survives casual testing because most tiles are fine.

Tuning It

Batch size: throughput and memory Throughput rises steeply from a batch of one to about sixteen, then flattens. Device memory rises linearly throughout. The useful operating point is the knee, where throughput has stopped improving and memory headroom remains for the activations. Throughput flattens; memory does not the knee — batch 16 tiles per second device memory 1 8 16 48 128 Batch size. Find the knee once per model and tile size; it barely moves after that.

Three measurements tell you what to change. Device utilisation below 70% means the readers are behind — add threads, or check whether the source is a striped GeoTIFF forcing full-width reads. A queue that is persistently full means the readers are ahead and the device is the limit, so a larger batch or a smaller model is the lever. And host memory climbing steadily means the queue is unbounded or predictions are being accumulated rather than written.

Half precision is worth trying once the pipeline is otherwise balanced. Casting the model and inputs to float16 typically halves the model time on hardware with tensor cores, which matters only if the device was the bottleneck — on a read-bound pipeline it changes nothing, which is itself a useful diagnostic.

Common Errors

Occasional tiles come back as noise

A dataset handle was shared between threads. Open the file inside each worker, as the example does.

RuntimeError: CUDA out of memory partway through

The batch is too large for the activations, or a variable-shape edge batch was larger than expected. Pad edge tiles to the nominal size and keep the batch fixed.

Throughput is the same with and without threads

The source is a striped GeoTIFF, so every window read pulls full-width rows and the storage saturates immediately. Convert it with the guidance in converting a GeoTIFF to a COG with rio-cogeo.

Frequently Asked Questions

Q: Why threads rather than processes for the readers? Because GDAL releases the GIL while decompressing, so threads genuinely overlap, and because the arrays stay in one address space — a process pool would have to pickle every tile back to the parent, which costs more than the read.

Q: How many reader threads should I use? Two to four for local NVMe, eight to sixteen for object storage where latency rather than bandwidth is the limit. Measure the queue depth: if it is usually full the readers are ahead, and if it is usually empty add more.

Q: Does a bigger batch always help? Only until the device is saturated, which for a small segmentation network is usually between eight and thirty-two tiles. Beyond that the throughput curve is flat and the memory cost is real, so stop at the knee.

Q: Does this pattern help on CPU-only machines? Yes, though less dramatically. A CPU model competes with the readers for cores, so the gain comes from overlapping decompression with compute rather than from filling an idle device. Use one or two readers rather than four, and set OMP_NUM_THREADS so the model does not oversubscribe the machine and starve the readers it depends on.