Train a model on your dataset

This tutorial trains a small U-Net on a mapcv dataset in plain PyTorch. It works unchanged on the RGB building dataset from Buildings from aerial imagery and on the four-band dataset from Land cover from Sentinel-2. The goal is a correct, honest baseline, not a state-of-the-art score.

You need: a mapcv dataset with labels and splits, and pip install torch numpy pillow. A GPU helps but is not required for these small datasets.

A notebook version is in the repository: examples/notebooks/02-train-a-segmentation-model.ipynb.

1. Load the splits

   mapcv.data.MapcvDataset reads one split. With normalize=True it standardises each band with the mean and standard deviation of the train split (read from stats.json, or computed in memory when it is missing; run mapcv stats once to save them), so RGB and Sentinel-2 inputs need no further scaling, and pixels without imagery become 0. Each item is a dict with an image tensor (C, H, W) and a mask tensor (H, W); ignore_index is the mask value mapcv writes where there is no imagery (255).

   import torch
   from torch import nn
   from torch.utils.data import DataLoader
   
   from mapcv.data import MapcvDataset
   
   root = "dataset"
   train_ds = MapcvDataset(root, split="train", normalize=True)
   val_ds = MapcvDataset(root, split="val", normalize=True)
   test_ds = MapcvDataset(root, split="test", normalize=True)
   
   IGNORE_INDEX = train_ds.ignore_index                    # 255
   in_channels = train_ds[0]["image"].shape[0]             # 3 for RGB, 4 for b04/b03/b02/b08
   num_classes = max(train_ds.classes, default=1) + 1      # + background
   print(len(train_ds), len(val_ds), len(test_ds), in_channels, num_classes)
   
   
   def batches(dataset, **kwargs):
       """(image, mask) batches: float32 images, int64 masks for the loss."""
       for batch in DataLoader(dataset, **kwargs):
           yield batch["image"].float(), batch["mask"].long()

2. Define a small U-Net

   def block(c_in, c_out):
       return nn.Sequential(
           nn.Conv2d(c_in, c_out, 3, padding=1), nn.BatchNorm2d(c_out), nn.ReLU(inplace=True),
           nn.Conv2d(c_out, c_out, 3, padding=1), nn.BatchNorm2d(c_out), nn.ReLU(inplace=True),
       )
   
   
   class TinyUNet(nn.Module):
       def __init__(self, in_channels, num_classes, width=32):
           super().__init__()
           self.enc1 = block(in_channels, width)
           self.enc2 = block(width, width * 2)
           self.enc3 = block(width * 2, width * 4)
           self.pool = nn.MaxPool2d(2)
           self.up2 = nn.ConvTranspose2d(width * 4, width * 2, 2, stride=2)
           self.dec2 = block(width * 4, width * 2)
           self.up1 = nn.ConvTranspose2d(width * 2, width, 2, stride=2)
           self.dec1 = block(width * 2, width)
           self.head = nn.Conv2d(width, num_classes, 1)
   
       def forward(self, x):
           e1 = self.enc1(x)
           e2 = self.enc2(self.pool(e1))
           e3 = self.enc3(self.pool(e2))
           d2 = self.dec2(torch.cat([self.up2(e3), e2], dim=1))
           d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
           return self.head(d1)

   Patch sizes must be divisible by 4 for this network; 128 and 256 are. For real projects, a pretrained encoder (for example from segmentation-models-pytorch) will do much better, especially on small datasets.

3. Measure IoU

   Intersection-over-union per class, from a confusion matrix that skips ignored pixels:

   @torch.no_grad()
   def per_class_iou(model, dataset, device):
       confusion = torch.zeros(num_classes, num_classes, dtype=torch.long)
       model.eval()
       for x, y in batches(dataset, batch_size=8):
           pred = model(x.to(device)).argmax(dim=1).cpu()
           keep = y != IGNORE_INDEX
           confusion += torch.bincount(
               y[keep] * num_classes + pred[keep], minlength=num_classes**2
           ).reshape(num_classes, num_classes)
       intersection = confusion.diag().float()
       union = confusion.sum(0).float() + confusion.sum(1).float() - intersection
       return intersection / union.clamp(min=1)

4. Train

   device = "cuda" if torch.cuda.is_available() else "cpu"
   model = TinyUNet(in_channels, num_classes).to(device)
   optimiser = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
   loss_fn = nn.CrossEntropyLoss(ignore_index=IGNORE_INDEX)
   
   best = 0.0
   for epoch in range(30):
       model.train()
       for x, y in batches(train_ds, batch_size=8, shuffle=True, num_workers=2):
           # Flips and 90° rotations are label-preserving for overhead imagery.
           if torch.rand(1) < 0.5:
               x, y = x.flip(-1), y.flip(-1)
           k = int(torch.randint(4, (1,)))
           x, y = torch.rot90(x, k, (-2, -1)), torch.rot90(y, k, (-2, -1))
   
           x, y = x.to(device), y.to(device)
           loss = loss_fn(model(x), y)
           optimiser.zero_grad()
           loss.backward()
           optimiser.step()
   
       val_iou = per_class_iou(model, val_ds, device)
       print(f"epoch {epoch:2d}  loss {loss.item():.3f}  val mIoU {val_iou.mean():.3f}")
       if val_iou.mean() > best:
           best = val_iou.mean().item()
           torch.save(model.state_dict(), "best.pt")

   Use the val split to choose the checkpoint, learning rate and epochs. With rare classes (buildings were 19 % of pixels), consider class weights from the manifest: see Use your dataset → class weights, then pass weight= to CrossEntropyLoss.

5. Test once

   model.load_state_dict(torch.load("best.pt"))
   test_iou = per_class_iou(model, test_ds, device)
   for cid, value in enumerate(test_iou.tolist()):
       print(cid, f"{value:.3f}")
   print("test mIoU", test_iou.mean().item())

   Report this number, computed once, after all choices were made on val. Because the default split is spatial, the test patches come from blocks of the map the model never saw, so the score is a fair estimate of how it does on new areas nearby (see Spatial leakage).

What to expect

On the 77-patch dataset from the quickstart (RGB patches, PyTorch on the CPU of a 4-core machine), step 1 prints the split sizes 52 8 17 first, and the training loop prints one epoch … loss … val mIoU … line per epoch.

• Time: about one minute per epoch (3 epochs took 2 min 54 s), so the 30 epochs of step 4 take about half an hour on a CPU. A GPU is much faster.
• Score: after 3 epochs the val and test mIoU were both about 0.41, and the IoU of the building class about 0.06, so the model still predicts mostly background. Treat that as a sign that loading, training and scoring work together, not as a result: with 8 val and 17 test patches the numbers move with the split seed, and they improve with more epochs, more patches and a pretrained encoder.
• Check the data first if step 1 prints a split with 0 patches (re-run mapcv split and read its warnings) or if the loss is nan.

Tiny datasets give noisy scores

The tutorial datasets have a few dozen patches and only a handful in val and test. They are good for checking that the pipeline works, not for comparing models. Generate a larger region (see Large regions and resuming) before drawing conclusions, and try several split seeds (mapcv split dataset --seed 1, --seed 2, …).

Going further

Semi-supervised trainingUse splits//labeled.txt and unlabeled.txt for label-budget experiments.

Map predictions backTurn patch predictions into georeferenced outputs with the manifest transform.