U-Net

U-Net is a convolutional neural network (CNN) for semantic segmentation, the task of assigning a class to every pixel of an image. Olaf Ronneberger, Philipp Fischer, and Thomas Brox of the University of Freiburg introduced it in 2015 in U-Net: Convolutional Networks for Biomedical Image Segmentation. An encoder shrinks the image to capture what is in it, a decoder expands the result back to full resolution, and skip connections hand the decoder the encoder's fine detail at every scale, so object boundaries come out sharp. U-Net trained well on a few dozen annotated microscopy images, won two categories of the ISBI 2015 cell tracking challenge by a large margin, and remains a standard model in medical image analysis. Its design later became the denoising network of many diffusion models.

How U-Net Works#

The network has two symmetric halves, drawn as the two arms of a U:

• Contracting path (encoder): Repeated pairs of 3×3 convolutions with ReLU activations, each pair followed by 2×2 max pooling that halves the resolution. The number of channels doubles at every level, from 64 to 1,024, so the network trades spatial detail for a wider receptive field and richer features.
• Expanding path (decoder): At each level, a 2×2 up-convolution doubles the resolution and halves the channels, the result is concatenated with the encoder's feature map from the same level, and two 3×3 convolutions merge them. A final 1×1 convolution maps each pixel's features to class scores.

The original network has 23 convolutional layers and four pooling steps. It used unpadded convolutions, so a 572×572 input tile produced a 388×388 segmentation map, and large images were segmented as overlapping tiles. Most implementations today pad the convolutions so the output matches the input size.

Skip Connections#

Pooling discards exactly where each feature was. By the bottleneck, the network knows that a region contains a cell membrane but not which pixels the membrane passes through. The skip connections copy each encoder level's high-resolution feature map across the U to the decoder, which combines "what" from the upsampled path with "where" from the encoder. U-Net concatenates the two along the channel dimension, while ResNet residual connections add features inside a block. This PyTorch code builds one encoder level, a bottleneck, and one decoder level:

import torch
from torch import nn

# One encoder level, a bottleneck, and one decoder level of a U-Net
encoder = nn.Sequential(nn.Conv2d(3, 16, 3, padding=1), nn.ReLU())
bottleneck = nn.Sequential(nn.MaxPool2d(2), nn.Conv2d(16, 32, 3, padding=1), nn.ReLU())
upsample = nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2)
decoder = nn.Sequential(nn.Conv2d(32, 16, 3, padding=1), nn.ReLU())
head = nn.Conv2d(16, 2, kernel_size=1)  # 2 classes, e.g. cell and background

image = torch.randn(1, 3, 128, 128)
skip = encoder(image)  # full-resolution features kept for the decoder
context = bottleneck(skip)  # half resolution, twice the channels, wider context

# Skip connection: concatenate upsampled context with the saved high-resolution features
merged = torch.cat([upsample(context), skip], dim=1)
logits = head(decoder(merged))
print(f"context {tuple(context.shape)} -> merged {tuple(merged.shape)} -> logits {tuple(logits.shape)}")

It prints context (1, 32, 64, 64) -> merged (1, 32, 128, 128) -> logits (1, 2, 128, 128): the 16 upsampled context channels and the 16 skip channels stack into 32, and the output holds two class scores for each of the input's 128×128 pixels.

Training With Little Data#

The original training set for neuronal structures in electron microscopy contained 30 images of 512×512 pixels. U-Net made up for it with heavy data augmentation, especially random elastic deformations that mimic how tissue bends, and with a weighted cross-entropy loss that emphasizes the thin background gaps between touching cells, so neighboring cells come out as separate regions. Medical segmentation models now often train with the Dice loss, which measures mask overlap directly and copes with a small structure occupying only a few percent of the pixels.

Variants and Descendants#

• 3D U-Net and V-Net: Replace 2D operations with 3D ones to segment volumes such as CT and MRI scans.
• Attention U-Net and UNet++: Add attention gates or nested, densely connected skip paths to refine which encoder features reach the decoder.
• nnU-Net: A framework that configures preprocessing, architecture, and training for a new dataset automatically. A MICCAI 2024 study found that well-tuned CNN-based U-Nets in this framework still outperform newer transformer and Mamba architectures on 3D medical segmentation when validated rigorously.
• Diffusion U-Nets: DDPM and Stable Diffusion 1.x use a U-Net to predict the noise in an image or latent, with attention layers added so the prompt can steer generation; Stable Diffusion 1.x's U-Net has about 860 million parameters, and ControlNet attaches to this U-Net's encoder. Newer generators replace the U-Net with diffusion transformers.

U-Net vs. Related Models#

• Fully convolutional networks (FCN): The 2014 FCN turned classifiers into per-pixel predictors and fused coarse and fine predictions by adding class score maps. U-Net built a full decoder with feature concatenation at every level.
• Autoencoders: Also shaped as an encoder and a decoder, but trained to reconstruct their input through a narrow bottleneck. Skip connections would let an autoencoder bypass that bottleneck, while U-Net needs them to produce aligned per-pixel labels.
• Instance segmentation models: A U-Net outputs one class per pixel, so two touching objects of one class share a label unless training or post-processing separates them.

Ultralytics YOLO26 semantic segmentation models, marked by the -sem suffix, use a shortened form of the U-Net decoder. Their neck upsamples the 1/32-scale backbone output twice, concatenating backbone features at 1/16 and 1/8 scale along the way, then classifies pixels at 1/8 scale and upsamples the class scores 8× to the input size.

Real-World Applications#

• Satellite image analysis: Land-cover maps label every pixel as water, forest, cropland, or buildings to track deforestation and urban growth.
• Industrial inspection: Pixel masks of scratches, cracks, and corrosion measure a defect's exact size and shape, which a bounding box cannot.
• Medical imaging: U-Nets outline organs, tumors, and lesions in CT and MRI scans for diagnosis and radiotherapy planning, and separate individual cells in microscopy. In the ISBI 2015 cell tracking challenge, U-Net trained on 20 partially annotated images of HeLa cells reached 77.5% average IoU, against 46% for the second-best method.