U-Net Image Segmentation

Six Segmentation Variants Measured on the Same Data

Computer vision · Oct 2026 · Completed

Which U-Net design choices actually matter: skip connections, pooling, block depth and residuals, each one a flag and a number.

Pet IoU 0.679 / Dice 0.809
Skips: IoU 0.489 → 0.560

Built withPyTorchtorchvisionpytestMatplotlibNumPyKaggleGoogle Colab

View code on GitHub

Overview

This started as a computer-vision coursework notebook that asked four questions about U-Net: does max-pooling beat a strided convolution, how much do skip connections bring, how does an encoder-decoder compare with a fully connected network, and how should the inference threshold be chosen. I rebuilt it as a package in which each question is one flag, the models are small (1 to 2M parameters, minutes per run on a CPU) and every answer is re-measured with IoU and Dice instead of pixel accuracy.

Six presets share one EncoderDecoder class: the plain U-Net, a strided-conv version, a single-conv version, a residual version, a SegNet-like model with no skip connections, and a fully connected baseline made of Linear layers only. They all train on the same synthetic shapes and on a subset of Oxford-IIIT Pet, with BCE + Dice loss, and are scored from one confusion matrix. A threshold sweep on the validation split picks the cutoff, never the test split.

On the shapes the plain U-Net reaches test IoU 0.560 against 0.422 for strided convolutions and 0.489 without skips; the fully connected model scores 0.102 with ten times the parameters. On the pet photos the residual U-Net leads at 0.679, then U-Net 0.641, no skips 0.630 and fully connected 0.524. Porting the notebook also exposed that its validation and test sets were augmented, its threshold sweep ran over 40 to 80 instead of 0.40 to 0.80, and its fully connected model ended in a ReLU that forbade negative logits.

The question

What do skip connections, max-pooling, double convolutions and residual blocks each buy a segmentation network, when every variant is trained on the same data and scored with the same metric?

Key features

Six presets from one model class: unet, unet-strided, unet-single, unet-residual, segnet (no skips) and fcn, plus switches to combine them freely.
Loss as a flag: per-pixel BCE, soft Dice, or both (the default), so the effect of an overlap loss on the probabilities can be measured.
Metrics from one confusion matrix: IoU, Dice, pixel accuracy, precision, recall and per-image mean IoU; checkpoints are selected on validation IoU.
A threshold sweep from 0.1 to 0.9 on the validation split, with the test score reported at 0.5 and at the chosen threshold.
Two datasets through one command: an offline synthetic shapes generator (800 / 100 / 100 images in 4 seconds) and a random subset of Oxford-IIIT Pet with binary masks.
Comparison and before/after scripts that train every variant on the same data, write the results table, and record predictions after every epoch into figures and GIFs.
29 pytest checks, including a test that the U-Net overfits one batch and that the fully connected head can output negative logits, plus a Kaggle / Colab notebook that explains each variant between the cells.

How it works

Down-sample four times
A stem ConvBlock followed by four down steps halves the resolution each time, by a fixed 2 × 2 max-pool or a learned 3 × 3 stride-2 convolution, with one or two convolutions per block and an optional residual path.
Up-sample and concatenate
Four ConvTranspose 2 × 2 steps bring the bottleneck back to full size; with skips on, each step concatenates the encoder map of the same size, and with skips off the decoder only sees the bottleneck.
Or skip the convolutions
The fully connected baseline flattens the image, passes it through a stack of Linear layers and reshapes the output, so it has no locality and no hierarchy at all.
Train with BCE + Dice
AdamW with warm-up and cosine decay, batch 16, 96 px images; BCE scores every pixel and soft Dice scores the overlap, which pushes the probabilities toward 0 or 1.
Score and choose the threshold
Probabilities above a threshold become the mask; IoU, Dice, precision and recall are read off one confusion matrix, and the threshold is swept on the validation split before the test split is scored.

Examples

U-Net on the synthetic shapes, one frame per epoch: green is correct object pixels, red false positives, blue missed pixels. The untrained network paints almost everything red; by the last epoch the mask follows the outlines.

The same U-Net on Oxford-IIIT Pet photos, epoch by epoch: the mask moves onto the animal within a few epochs, and this re-run reproduced the table's test IoU of 0.641 exactly.

The same test images through every variant, with the per-image IoU below each panel. The strided model misses whole small triangles, the no-skip model turns shapes into rounded blobs, and the fully connected model paints noise where objects tend to be.

Pet photos through U-Net, residual U-Net, the no-skip model and the fully connected baseline. The U-Nets find the animal (recall about 0.85) but also paint background of a similar colour; the fully connected model predicts a centred blob whatever the pose.

Validation IoU and training loss per epoch for the six variants on the shapes. The fully connected model's IoU falls after epoch 2 as it starts memorising the 600 training images.

Input

unetseg-data shapes && python scripts/compare_all.py --data data/shapes --epochs 6 --max-train 600 --threads 1 --out runs/shapes

Output

test IoU: unet-single 0.660, unet-residual 0.633, unet 0.560, segnet 0.489, unet-strided 0.422, fcn 0.102 (Dice 0.796, 0.775, 0.718, 0.657, 0.593, 0.184); pixel accuracy above 0.83 for every model, including fcn.

The run behind the shapes table: 600 training images at 96 px, 6 epochs, BCE + Dice, 2 to 11 minutes per model on one CPU thread.

Tech stack

PyTorch
The model variants, losses, metrics and training loop
torchvision
Oxford-IIIT Pet download and the default joint augmentations
pytest
29 checks on models, data splits, metrics and overfitting one batch
Matplotlib
Curves, prediction grids, TP / FP / FN overlays, threshold and confusion plots
NumPy
Shapes generator and confusion-matrix metrics
Kaggle
GPU notebook that trains a wider U-Net on all 7,390 Pet photos
Google Colab
The same notebook

What I learned

Pixel accuracy hides everything: every model on the shapes, including the fully connected one, scores above 0.83, so IoU is the number to select checkpoints on
Six-epoch runs are noisy: the same U-Net reached test IoU 0.47 or 0.56 depending only on the thread count, so differences under 0.1 between shapes rows are noise

Limitations

The numbers are short CPU runs (600 training images, 96 px, 6 epochs) meant to compare variants, not benchmarks; the full-dataset GPU run in the notebook has no reported numbers
The single-conv and residual U-Nets lead the plain U-Net after 6 epochs because they learn faster; the plain one was still climbing, so the final order is unknown

Keywords

U-Net · image segmentation · binary segmentation · skip connections · max-pooling vs strided convolution · residual blocks · SegNet · fully convolutional network · Dice loss · IoU · threshold sweep · Oxford-IIIT Pet · PyTorch · computer vision · deep learning