Building a U-Net for Land Cover Segmentation in PyTorch

Build and train a U-Net on multi-spectral patches in PyTorch: channel-flexible input, ignore-index loss, class weighting, AMP training loop and per-class IoU tracking.

The architecture is an encoder that halves resolution four times, a bottleneck, and a decoder that doubles it back, with skip connections carrying full-resolution detail across. What makes it work on satellite imagery rather than photographs is the surrounding detail: an input layer sized to the number of spectral bands, a loss that ignores unlabelled pixels, class weights that stop the majority class from swallowing the gradient, and a validation set that is geographically separate.

This page is the implementation. For the design decisions around it — splits, patch strategy, inference — see Deep Learning Segmentation for Satellite Imagery.

One training step, with the two geospatial-specific pieces marked A batch of multi-spectral patches is normalised using statistics loaded from the training artifact, passed through the U-Net to produce per-class logits, and compared with the label patch by a loss that excludes ignore-valued pixels. The resulting gradient updates the weights. The normalisation source and the ignore handling are highlighted as the two steps specific to geospatial data. The architecture is standard; the two highlighted steps are not patch batch B × 6 × 256 × 256 normalise stats LOADED from the training artifact U-Net → B × C × 256 × 256 weighted loss ignore_index = 255 + Dice term step next batch Recomputing the statistics per batch, or letting ignore pixels into the loss, breaks the model quietly.

Why This Fails in Geospatial ML Pipelines

The commonest failure is a loss that treats unlabelled pixels as a class. Label rasters carry an ignore value for ground nobody digitised, and unless it is passed to the loss as ignore_index, the network is trained to predict “class 255” — or, if the ignore value happens to equal the background code, to paint background over every undigitised region. Since undigitised ground is usually the majority of a scene, the model converges beautifully and produces a map that is uniformly one class.

The second is nodata reaching the tensors. Satellite products encode gaps with sentinels like -9999. After normalisation that becomes a number several hundred standard deviations from the mean, which saturates the batch norms and drives the loss to nan within a few hundred steps. The traceback points at the loss; the cause is three files upstream in the reader.

The third is a validation set that is not really held out. If patches were split randomly rather than by block, validation IoU rises smoothly, early stopping fires late, and the resulting model is tuned on data it memorised. The remedy is upstream, in tiling large rasters into training patches, but its symptom shows up here as a validation curve that looks too good.

Core Principles

  • Size the first convolution from the band count, never assume three channels.
  • Pass ignore_index to the loss and confirm it with a test.
  • Cap class weights. Inverse frequency unbounded makes a 0.1% class destabilise training.
  • Track macro IoU, not accuracy. Stop on the metric you actually care about.
  • Use mixed precision. It buys a larger patch, which buys more context per prediction.
  • Checkpoint on the metric, not the loss. They diverge, and IoU is the one that matters.

Production-Ready Code

from __future__ import annotations

import logging
from dataclasses import dataclass

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset

logger = logging.getLogger(__name__)
IGNORE_INDEX = 255


def conv_block(cin: int, cout: int) -> nn.Sequential:
    return nn.Sequential(
        nn.Conv2d(cin, cout, 3, padding=1, bias=False),
        nn.BatchNorm2d(cout), nn.ReLU(inplace=True),
        nn.Conv2d(cout, cout, 3, padding=1, bias=False),
        nn.BatchNorm2d(cout), nn.ReLU(inplace=True),
    )


class UNet(nn.Module):
    """U-Net whose input width is set by the number of spectral bands."""

    def __init__(self, in_bands: int, n_classes: int, width: int = 32):
        super().__init__()
        w = width
        self.enc = nn.ModuleList([
            conv_block(in_bands, w), conv_block(w, w * 2),
            conv_block(w * 2, w * 4), conv_block(w * 4, w * 8),
        ])
        self.pool = nn.MaxPool2d(2)
        self.bottleneck = conv_block(w * 8, w * 16)
        self.ups = nn.ModuleList([
            nn.ConvTranspose2d(w * 16, w * 8, 2, stride=2),
            nn.ConvTranspose2d(w * 8, w * 4, 2, stride=2),
            nn.ConvTranspose2d(w * 4, w * 2, 2, stride=2),
            nn.ConvTranspose2d(w * 2, w, 2, stride=2),
        ])
        self.dec = nn.ModuleList([
            conv_block(w * 16, w * 8), conv_block(w * 8, w * 4),
            conv_block(w * 4, w * 2), conv_block(w * 2, w),
        ])
        self.head = nn.Conv2d(w, n_classes, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        skips = []
        for enc in self.enc:
            x = enc(x)
            skips.append(x)
            x = self.pool(x)
        x = self.bottleneck(x)
        for up, dec, skip in zip(self.ups, self.dec, reversed(skips)):
            x = dec(torch.cat([up(x), skip], dim=1))
        return self.head(x)


class PatchDataset(Dataset):
    """Memory-mapped patch store, normalised with statistics from the artifact."""

    def __init__(self, x_path: str, y_path: str, indices: list[int],
                 mean: np.ndarray, std: np.ndarray, augment: bool = False):
        self.x = np.load(x_path, mmap_mode="r")
        self.y = np.load(y_path, mmap_mode="r")
        self.idx = indices
        self.mean = mean[:, None, None].astype("float32")
        self.std = std[:, None, None].astype("float32")
        self.augment = augment

    def __len__(self) -> int:
        return len(self.idx)

    def __getitem__(self, i: int):
        j = self.idx[i]
        x = np.asarray(self.x[j], dtype="float32")
        y = np.asarray(self.y[j], dtype="int64")
        if not np.isfinite(x).all():
            x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
        x = (x - self.mean) / self.std

        if self.augment:                       # only flips/rotations — never colour jitter
            k = np.random.randint(4)
            x, y = np.rot90(x, k, (1, 2)).copy(), np.rot90(y, k, (0, 1)).copy()
            if np.random.rand() < 0.5:
                x, y = x[:, :, ::-1].copy(), y[:, ::-1].copy()
        return torch.from_numpy(x), torch.from_numpy(y)


def capped_class_weights(counts: np.ndarray, cap: float = 8.0) -> torch.Tensor:
    """Inverse-frequency weights, clipped so a very rare class cannot dominate."""
    w = counts.sum() / (len(counts) * np.maximum(counts, 1))
    return torch.tensor(np.clip(w, 0.2, cap), dtype=torch.float32)


def iou_from_confusion(cm: np.ndarray) -> np.ndarray:
    tp = np.diag(cm).astype("float64")
    denom = cm.sum(axis=1) + cm.sum(axis=0) - tp
    return np.where(denom > 0, tp / np.maximum(denom, 1), np.nan)

Step-by-Step Walkthrough

Step 1 — Wire the loaders to the manifest splits

import json

blob = json.load(open("data/patches_manifest.json"))
splits = {s: [p["index"] for p in blob["patches"] if p["split"] == s]
          for s in ("train", "val", "test")}
stats = np.load("artifacts/norm_stats.npz")

train_ds = PatchDataset("data/patches_x.npy", "data/patches_y.npy",
                        splits["train"], stats["mean"], stats["std"], augment=True)
val_ds = PatchDataset("data/patches_x.npy", "data/patches_y.npy",
                      splits["val"], stats["mean"], stats["std"])

train_dl = DataLoader(train_ds, batch_size=8, shuffle=True, num_workers=4,
                      pin_memory=True, drop_last=True)
val_dl = DataLoader(val_ds, batch_size=8, num_workers=2, pin_memory=True)
print(f"train {len(train_ds)} patches, val {len(val_ds)} patches")

Loading the normalisation statistics from a file rather than computing them here is the contract described in scaling features consistently between training and inference.

Step 2 — Build the model and the loss

N_CLASSES = 5
counts = np.array([4_100_000, 820_000, 96_000, 31_000, 210_000], dtype="float64")

device = "cuda" if torch.cuda.is_available() else "cpu"
model = UNet(in_bands=stats["mean"].shape[0], n_classes=N_CLASSES).to(device)
criterion = nn.CrossEntropyLoss(weight=capped_class_weights(counts).to(device),
                                ignore_index=IGNORE_INDEX)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=40)
scaler = torch.cuda.amp.GradScaler(enabled=device == "cuda")

Step 3 — The training loop

best_iou = -1.0
for epoch in range(40):
    model.train()
    running = 0.0
    for x, y in train_dl:
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        optimizer.zero_grad(set_to_none=True)
        with torch.autocast(device_type=device, enabled=device == "cuda"):
            loss = criterion(model(x), y)
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
        running += float(loss)

    model.eval()
    cm = np.zeros((N_CLASSES, N_CLASSES), dtype="int64")
    with torch.no_grad():
        for x, y in val_dl:
            pred = model(x.to(device)).argmax(1).cpu().numpy()
            t = y.numpy()
            keep = t != IGNORE_INDEX
            k = t[keep] * N_CLASSES + pred[keep]
            cm += np.bincount(k, minlength=N_CLASSES ** 2).reshape(N_CLASSES, N_CLASSES)

    ious = iou_from_confusion(cm)
    macro = float(np.nanmean(ious))
    scheduler.step()
    logger.info("epoch %2d  loss %.4f  macro IoU %.4f  per-class %s",
                epoch, running / len(train_dl), macro, np.round(ious, 3).tolist())

    if macro > best_iou:
        best_iou = macro
        torch.save({"state_dict": model.state_dict(), "macro_iou": macro,
                    "bands": int(stats["mean"].shape[0]), "n_classes": N_CLASSES},
                   "artifacts/unet_best.pt")

Step 4 — Read the per-class curve, not the loss

print({i: round(float(v), 3) for i, v in enumerate(ious)})
assert not np.isnan(ious).all(), "no class was ever predicted — check ignore_index"

A class stuck at zero IoU across many epochs means the sampler never shows it, not that its weight is too low. Fix it in the patch sampler rather than by raising the weight past the cap.

Accuracy plateaus early; macro IoU is still climbing Two curves are plotted over forty epochs. Overall pixel accuracy rises within a few epochs to a high plateau above ninety percent and stays flat. Macro intersection over union rises much more slowly and is still improving when accuracy has long stopped moving, so stopping on accuracy would end training far too early. Validation curves over 40 epochs 5152540 accuracy macro IoU stop on accuracy stop on macro IoU Checkpoint on the metric you will be judged by, or you ship the epoch-8 model.

Verification

import numpy as np
import torch
import torch.nn as nn


def test_output_shape_matches_input():
    model = UNet(in_bands=6, n_classes=5)
    out = model(torch.zeros(2, 6, 256, 256))
    assert out.shape == (2, 5, 256, 256)


def test_ignore_pixels_contribute_no_gradient():
    """A fully unlabelled batch must produce a zero or nan loss, never a real one."""
    logits = torch.randn(1, 5, 8, 8, requires_grad=True)
    target = torch.full((1, 8, 8), IGNORE_INDEX, dtype=torch.long)
    loss = nn.CrossEntropyLoss(ignore_index=IGNORE_INDEX)(logits, target)
    assert torch.isnan(loss) or float(loss) == 0.0


def test_model_overfits_a_single_patch():
    """The strongest smoke test there is: if it cannot memorise one patch, it is broken."""
    torch.manual_seed(0)
    model = UNet(in_bands=3, n_classes=2, width=8)
    x = torch.randn(1, 3, 64, 64)
    y = (torch.rand(1, 64, 64) > 0.5).long()
    opt = torch.optim.Adam(model.parameters(), lr=1e-2)
    crit = nn.CrossEntropyLoss()
    for _ in range(120):
        opt.zero_grad()
        loss = crit(model(x), y)
        loss.backward()
        opt.step()
    assert float(loss) < 0.05, f"could not overfit one patch (loss {float(loss):.3f})"


def test_augmentation_keeps_x_and_y_aligned():
    ds = PatchDataset("data/patches_x.npy", "data/patches_y.npy", [0],
                      mean=np.zeros(6, "float32"), std=np.ones(6, "float32"), augment=True)
    x, y = ds[0]
    assert x.shape[-2:] == y.shape[-2:]

The overfit test is worth more than every other check combined. A model that cannot drive the loss to near zero on a single patch has a wiring bug — a detached skip connection, a wrong channel count, a target that is not aligned with its input — and no amount of training will rescue it.

Geometry augmentations are safe; radiometric ones are not A patch and its label are rotated together by ninety degrees, showing a valid augmentation because the transformation is applied identically to both. Beside it, a brightness jitter is marked invalid because it alters calibrated reflectance values, which carry physical meaning the model relies on. Rotate and flip freely — do not jitter reflectance x y rot90 x′ y′ same k applied to both → valid x brightness x″ NDVI thresholds no longer hold reflectance is a calibrated physical quantity, not an aesthetic property of a photograph

FAQ

Can I use an ImageNet-pretrained encoder with six spectral bands?

Yes — replace the first convolution and initialise the new channels from the mean of the pretrained RGB weights. Early filters learn edges and textures that transfer, though the benefit is smaller than on natural images because the spectral statistics differ so much.

Why does my loss go to nan in the first epoch?

Almost always an unmasked nodata sentinel: -9999 survives normalisation as a huge negative value and saturates the batch norms. Convert nodata to NaN at read time and handle it before the loader, as the dataset class above does.

How many epochs does a segmentation model need?

Fewer than you would expect, because each patch supplies tens of thousands of labelled pixels — 30–60 epochs over a few thousand patches is typical. Stop on macro IoU measured against a spatially separate validation set, per spatial cross-validation strategies.

Part of: Deep Learning Segmentation for Satellite Imagery Part of: Training Geospatial Predictive Models in Python