U-Net

U-Net is a convolutional network for semantic segmentation introduced by Olaf Ronneberger, Philipp Fischer and Thomas Brox at MICCAI 2015. It is an encoder–decoder: a contracting path of convolutions and pooling gathers context, and an expanding path of the same depth brings the resolution back, receiving at every scale a copy of the contracting path’s features through a skip connection. Drawn with the contracting path on the left and the expanding path on the right, the network forms a U. It was built for microscopy, where a training set may hold 20 or 30 annotated images, and it remains a standard architecture in medical image segmentation and a common backbone wherever an image must be mapped to an image of the same size.

Problem Addressed

Biomedical segmentation needs a label for every pixel and rarely has thousands of training images. The best method on the ISBI 2012 electron-microscopy (EM) challenge before U-Net, Ciresan et al.’s network, classified each pixel from a patch around it. Ronneberger et al. name two drawbacks: it is slow, because the network runs once per pixel on heavily overlapping patches, and it trades localization against context, because larger patches need more max pooling, which blurs position.

The fully convolutional network (FCN) of Long, Shelhamer and Darrell had removed the per-patch redundancy by predicting all pixels in one pass and upsampling coarse predictions. Compared with FCN, U-Net differs in three ways: a decoder with as many feature channels as the encoder, so context can propagate to the high-resolution layers; skips that concatenate features instead of adding class scores; and a training strategy built around heavy data augmentation, so that very few labeled images suffice.

Architecture

The network has no fully connected layers and uses only the valid part of each convolution, without padding, so every output pixel has its full context inside the input. Feature maps lose a border at every convolution and the output is smaller than the input (Ronneberger et al., 2015, Fig. 1).

Contracting Path

Four steps, each two unpadded $3 \times 3$ convolutions followed by ReLU, then $2 \times 2$ max pooling with stride 2. The number of channels doubles at each downsampling: 64, 128, 256, 512, and 1,024 in the two convolutions at the bottom of the U, at 1/16 of the input resolution.

Expanding Path

Each of four steps upsamples the feature map and applies a $2 \times 2$ convolution, the “up-convolution”, which halves the number of channels. The result is concatenated with the contracting-path feature map of the same scale, cropped to the same size because the unpadded convolutions have shrunk the decoder’s maps more, and two $3 \times 3$ convolutions with ReLU follow. The paper does not say how the upsampling is computed; a transposed convolution with a $2 \times 2$ kernel and stride 2 is one way to implement the step, and nearest-neighbor upsampling followed by a $2 \times 2$ convolution is another. A final $1 \times 1$ convolution maps each 64-channel feature vector to the class scores. The authors count 23 convolutional layers: by our count, 18 of size $3 \times 3$, the four up-convolutions and the $1 \times 1$ layer.

Tile Sizes and the Overlap-Tile Strategy

The published configuration takes a $572 \times 572$ input tile and outputs a $388 \times 388$ map. The authors require every $2 \times 2$ pooling to see an even size, so that tiles fit together. The arithmetic is quick to check:

import numpy as np

def unet_sizes(s, depth=4):
    """Feature-map side lengths through U-Net with unpadded 3x3 convolutions.
    Returns the list of sizes, or None if a 2x2 pooling meets an odd size."""
    sizes = [s]
    for _ in range(depth):                 # contracting path
        s -= 4                             # two 3x3 valid convolutions
        if s <= 0 or s % 2:
            return None                    # pooling would drop a row/column
        s //= 2                            # 2x2 max pooling, stride 2
        sizes.append(s)
    s -= 4                                 # two convolutions at the bottom
    if s <= 0:
        return None
    sizes.append(s)
    for _ in range(depth):                 # expansive path
        s = 2 * s - 4                      # 2x up-convolution, then two 3x3 convs
        if s <= 0:
            return None
        sizes.append(s)
    return sizes

print("572 ->", unet_sizes(572))

valid = [s for s in range(100, 1300) if unet_sizes(s)]
print("valid tile sizes:", valid[:6], "...", valid[-2:])
print("all of the form 16b + 60:", all((s - 60) % 16 == 0 for s in valid))
print("input - output is always", {s - unet_sizes(s)[-1] for s in valid})

# Overlap-tile: a 1024 x 1024 image with 572 x 572 input tiles.
H = W = 1024
tile_in = 572
tile_out = unet_sizes(tile_in)[-1]          # 388
border = (tile_in - tile_out) // 2          # 92 pixels of context per side
n = int(np.ceil(H / tile_out))              # tiles per axis
img = np.random.default_rng(0).random((H, W))
need = n * tile_out + 2 * border            # input extent the tiles read
padded = np.pad(img, ((border, need - H - border),) * 2, mode="reflect")
out = np.zeros((n * tile_out,) * 2)
for i in range(n):
    for j in range(n):
        tile = padded[i * tile_out:i * tile_out + tile_in, j * tile_out:j * tile_out + tile_in]
        assert tile.shape == (tile_in, tile_in)
        # stand-in for the network, to check the bookkeeping: the tile's centre
        out[i * tile_out:(i + 1) * tile_out, j * tile_out:(j + 1) * tile_out] = \
            tile[border:border + tile_out, border:border + tile_out]
print(f"{n} x {n} = {n * n} tiles; mirrored input {padded.shape}; "
      f"stitched output cropped back to the image: {np.array_equal(out[:H, :W], img)}")
print(f"input pixels read per output pixel: {n * n * tile_in ** 2 / (H * W):.2f}")
b = int(np.ceil((H + 124) / 16))            # smallest one-tile size with output >= 1024
print("one tile covering the image:", 16 * b + 60, "->", unet_sizes(16 * b + 60)[-1])

572 -> [572, 284, 140, 68, 32, 28, 52, 100, 196, 388]
valid tile sizes: [188, 204, 220, 236, 252, 268] ... [1276, 1292]
all of the form 16b + 60: True
input - output is always {184}
3 x 3 = 9 tiles; mirrored input (1348, 1348); stitched output cropped back to the image: True
input pixels read per output pixel: 2.81
one tile covering the image: 1212 -> 1028

Valid inputs are $16b + 60$ pixels wide, where $b$ is the size reaching the bottom of the U, and every one loses 184 pixels, 92 per side, so a $W$-pixel output tile needs a $(W + 184)$-pixel input. Ronneberger et al. segment a large image by this overlap-tile strategy: output tiles are laid side by side, their input tiles overlap, and context missing beyond the image border is filled by mirroring the image. A $1{,}024 \times 1{,}024$ image takes nine tiles of $572 \times 572$, which read 2.8 input pixels per output pixel; a single $1{,}212 \times 1{,}212$ tile would do, memory permitting. Without tiling, the authors note, GPU memory would limit the resolution of the images that can be processed.

Training

The authors trained with stochastic gradient descent in Caffe. They preferred large tiles to large batches, so a batch held one image, and set the momentum to 0.99 so that many previously seen samples contribute to each update (stochastic gradient descent). Weights were initialized from a Gaussian with standard deviation $\sqrt{2/N}$, $N$ being the number of incoming connections of a unit, which they call extremely important in a deep network with several paths.

Weighted Cross-Entropy

A pixel-wise softmax over the $K$ output channels gives probabilities $p_k(\mathbf{x})$, and the loss is a weighted cross-entropy summed over the output pixels $\Omega$:

$$E = -\sum_{\mathbf{x} \in \Omega} w(\mathbf{x}) \log p_{\ell(\mathbf{x})}(\mathbf{x}),$$

where $\ell(\mathbf{x})$ is the true label. (Eq. 1 of the arXiv version is printed without the minus sign.) The weight map is precomputed from each ground-truth mask:

$$w(\mathbf{x}) = w_c(\mathbf{x}) + w_0 \exp\!\left(-\frac{\big(d_1(\mathbf{x}) + d_2(\mathbf{x})\big)^2}{2\sigma^2}\right),$$

with $w_c$ balancing class frequencies, $d_1$ and $d_2$ the distances to the border of the nearest and second-nearest cell, $w_0 = 10$ and $\sigma \approx 5$ pixels. The second term matters for instance segmentation of cells. A per-pixel classifier sees two touching cells as one region; the training masks therefore contain a thin background gap between touching cells, made with morphological operations, and the weight map makes those few gap pixels expensive to misclassify. The script below computes the map for five synthetic cells. The class-balancing term is our choice, and the border term is applied only on background pixels, our reading of what the formula is for:

import numpy as np
from scipy import ndimage as ndi

# Five disk-shaped "cells" on a 64 x 64 image: 1-2 and 3-4 overlap, 5 is 3 px from 4.
H = W = 64
yy, xx = np.mgrid[:H, :W]
cells = [(20, 18, 9), (20, 34, 9), (46, 20, 8), (46, 35, 8), (47, 51, 5)]
inst = np.zeros((H, W), int)
for k, (cy, cx, r) in enumerate(cells, start=1):
    inst[((yy - cy) ** 2 + (xx - cx) ** 2 <= r ** 2) & (inst == 0)] = k

# Separation border: pixels of one cell adjacent to a different cell become
# background, so touching cells are split by a 1-2 pixel gap in the mask.
gap = np.zeros((H, W), bool)
for k in range(1, len(cells) + 1):
    near_other = ndi.binary_dilation((inst > 0) & (inst != k))
    gap |= (inst == k) & near_other
inst[gap] = 0
fg = inst > 0

# d1, d2: distance from each pixel to the nearest and second-nearest cell.
dist = np.stack([ndi.distance_transform_edt(inst != k) for k in range(1, len(cells) + 1)])
dist.sort(axis=0)
d1, d2 = dist[0], dist[1]

# w_c balances the two classes (our choice: weight inversely proportional to frequency);
# the border term follows Eq. 2 of Ronneberger et al. (2015), w0 = 10, sigma = 5.
w0, sigma = 10.0, 5.0
wc = np.where(fg, 0.5 / fg.mean(), 0.5 / (~fg).mean())
w = wc + np.where(fg, 0.0, w0 * np.exp(-(d1 + d2) ** 2 / (2 * sigma ** 2)))

narrow = ~fg & ~gap & (d1 + d2 <= 6)                     # other thin gaps between cells
rest = ~fg & ~gap & ~narrow
print(f"foreground fraction {fg.mean():.3f}; w_c = {wc[fg][0]:.2f} (cells), {wc[~fg][0]:.2f} (background)")
for name, m in [("carved gaps between touching cells", gap), ("other background, d1 + d2 <= 6", narrow),
                ("all other background", rest), ("cell pixels", fg)]:
    print(f"{name:35s} {m.sum():5d} px, w from {w[m].min():5.2f} to {w[m].max():5.2f}, mean {w[m].mean():5.2f}")
print(f"share of total loss weight on the gaps: {w[gap].sum() / w.sum():.3f} "
      f"(they are {gap.mean():.3f} of the pixels)")

foreground fraction 0.229; w_c = 2.19 (cells), 0.65 (background)
carved gaps between touching cells     32 px, w from  9.00 to  9.55, mean  9.14
other background, d1 + d2 <= 6         98 px, w from  5.52 to  9.88, mean  7.18
all other background                 3030 px, w from  0.65 to  5.40, mean  0.94
cell pixels                           936 px, w from  2.19 to  2.19, mean  2.19
share of total loss weight on the gaps: 0.050 (they are 0.008 of the pixels)

The 32 carved gap pixels weigh about 14 times as much as background far from any cell (0.65) and carry 5% of the total loss weight from 0.8% of the pixels. Narrow natural gaps, such as the 3-pixel space between cells 4 and 5, are upweighted too, since $d_1 + d_2$ is small there as well.

Data Augmentation

With so few images, augmentation is what the authors call the key to training. Microscopy needs invariance to shift and rotation and robustness to deformation and gray-value changes. Random elastic deformations do most of the work: displacement vectors drawn from a Gaussian with a standard deviation of 10 pixels on a coarse $3 \times 3$ grid, interpolated bicubically to a smooth per-pixel displacement field, by which image and mask are warped together. Dropout layers at the end of the contracting path add, in the paper’s words, “further implicit data augmentation”. Training took 10 hours on an NVIDIA Titan GPU with 6 GB.

Example: Border Weights on Touching Cells

The PyTorch code below (run with PyTorch 2.14.1 on a CPU in about 20 s, with SciPy for distance transforms and connected components) trains a two-level U-Net with padded convolutions and 8 to 32 channels on $48 \times 48$ images of five to eight noisy, dome-shaped cells, many overlapping. Targets carry the carved gaps. The same network is trained twice for 600 steps, with plain cross-entropy and with the border term of the weight map ($w_c = 1$ in both, so only the border term differs). Each prediction is thresholded and split into connected components, and a true cell counts as recovered when a component matches it with IoU above 0.5.

import sys
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
from scipy import ndimage as ndi

torch.set_num_threads(1)
SEED = int(sys.argv[1]) if len(sys.argv) > 1 else 8
S = 48

def make_image(rng):
    """Noisy 48x48 image of 5-8 dome-shaped cells, many of them touching.
    Returns the image, the instance labels, and the U-Net-style target mask
    with a background gap carved between touching cells."""
    yy, xx = np.mgrid[:S, :S]
    cells = []
    while len(cells) < rng.integers(5, 9):
        cy, cx, r = rng.uniform(6, S - 6), rng.uniform(6, S - 6), rng.uniform(4.5, 7)
        if all(np.hypot(cy - y, cx - x) > 0.6 * (r + q) for y, x, q in cells):
            cells.append((cy, cx, r))
    d = np.stack([np.hypot(yy - cy, xx - cx) / r for cy, cx, r in cells])  # 1 on each rim
    inside = d <= 1
    inst = np.where(inside.any(0), d.argmin(0) + 1, 0)       # overlaps go to the nearer centre
    img = np.where(inside, 0.4 + 0.6 * np.sqrt(np.clip(1 - d ** 2, 0, 1)), 0).max(0)
    img = img + 0.25 * rng.standard_normal((S, S))
    gap = np.zeros((S, S), bool)
    for k in range(1, len(cells) + 1):
        gap |= (inst == k) & ndi.binary_dilation((inst > 0) & (inst != k))
    return img, inst, (inst > 0) & ~gap

def border_weight(mask, w0=10.0, sigma=5.0):
    """w0 * exp(-(d1 + d2)^2 / (2 sigma^2)) on background pixels (Ronneberger et al., Eq. 2)."""
    lab, n = ndi.label(mask)
    if n < 2:
        return np.zeros(mask.shape)
    dist = np.sort(np.stack([ndi.distance_transform_edt(lab != k) for k in range(1, n + 1)]), 0)
    return np.where(mask, 0.0, w0 * np.exp(-(dist[0] + dist[1]) ** 2 / (2 * sigma ** 2)))

def dataset(n, seed):
    rng = np.random.default_rng(seed)
    imgs, insts, masks = zip(*(make_image(rng) for _ in range(n)))
    w = np.stack([border_weight(m) for m in masks])
    as_t = lambda a, dt: torch.tensor(np.stack(a), dtype=dt)
    return (as_t(imgs, torch.float32)[:, None], np.stack(insts),
            as_t(masks, torch.long), torch.tensor(w, dtype=torch.float32))

def double_conv(cin, cout):                     # padded 3x3 convolutions, unlike the original
    return nn.Sequential(nn.Conv2d(cin, cout, 3, padding=1), nn.ReLU(),
                         nn.Conv2d(cout, cout, 3, padding=1), nn.ReLU())

def up_conv(cin, cout):                         # upsampling, then a 2x2 convolution
    return nn.Sequential(nn.Upsample(scale_factor=2), nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(cin, cout, 2))

class TinyUNet(nn.Module):
    def __init__(self, c=16):
        super().__init__()
        self.enc0, self.enc1, self.mid = double_conv(1, c), double_conv(c, 2 * c), double_conv(2 * c, 4 * c)
        self.up1, self.dec1 = up_conv(4 * c, 2 * c), double_conv(4 * c, 2 * c)
        self.up0, self.dec0 = up_conv(2 * c, c), double_conv(2 * c, c)
        self.head = nn.Conv2d(c, 2, 1)          # background / cell scores per pixel
        for mod in self.modules():              # Gaussian, std sqrt(2 / fan-in), as in the paper
            if isinstance(mod, nn.Conv2d):
                nn.init.kaiming_normal_(mod.weight, nonlinearity="relu")
                nn.init.zeros_(mod.bias)

    def forward(self, x):
        e0 = self.enc0(x)
        e1 = self.enc1(F.max_pool2d(e0, 2))
        m = self.mid(F.max_pool2d(e1, 2))
        d1 = self.dec1(torch.cat([self.up1(m), e1], 1))     # skip: concatenate
        d0 = self.dec0(torch.cat([self.up0(d1), e0], 1))
        return self.head(d0)

def matched_cells(prob, inst):
    """Count true cells matched by a predicted connected component with IoU > 0.5."""
    lab, _ = ndi.label(prob > 0.5)
    hits = 0
    for k in range(1, inst.max() + 1):
        cell = inst == k
        ids, counts = np.unique(lab[cell], return_counts=True)
        counts, ids = counts[ids > 0], ids[ids > 0]
        if len(ids):
            comp = lab == ids[counts.argmax()]
            hits += (comp & cell).sum() / (comp | cell).sum() > 0.5
    return hits

x_tr, _, y_tr, wb_tr = dataset(400, seed=100 + SEED)
x_te, inst_te, y_te, _ = dataset(200, seed=999)                 # same test set for every run
n_cells = sum(i.max() for i in inst_te)
touching = sum(len(np.unique(i)) - 1 - ndi.label(i > 0)[1] for i in inst_te)

for use_border in (False, True):
    torch.manual_seed(SEED)
    net = TinyUNet(c=8)
    opt = torch.optim.Adam(net.parameters(), lr=2e-3)
    g = torch.Generator().manual_seed(SEED)
    for step in range(600):
        idx = torch.randint(0, len(x_tr), (8,), generator=g)
        ce = F.cross_entropy(net(x_tr[idx]), y_tr[idx], reduction="none")
        w = 1.0 + (wb_tr[idx] if use_border else 0.0)          # w_c = 1 for both runs
        loss = (w * ce).mean()
        opt.zero_grad()
        loss.backward()
        opt.step()
    with torch.no_grad():
        prob = net(x_te).softmax(1)[:, 1].numpy()
    pix_acc = ((prob > 0.5) == y_te.numpy().astype(bool)).mean()
    found = sum(matched_cells(p, i) for p, i in zip(prob, inst_te))
    comps = sum(ndi.label(p > 0.5)[1] for p in prob)
    print(f"{'with border weights' if use_border else 'plain cross-entropy'}: pixel accuracy {pix_acc:.3f}, "
          f"{comps} components, cells recovered {found}/{n_cells} ({found / n_cells:.1%})")
print(f"test set: {len(inst_te)} images, {n_cells} cells, {touching} merges if cells were not separated")

Output (seed 8):

plain cross-entropy: pixel accuracy 0.985, 1023 components, cells recovered 972/1167 (83.3%)
with border weights: pixel accuracy 0.979, 1154 components, cells recovered 1141/1167 (97.8%)
test set: 200 images, 1167 cells, 506 merges if cells were not separated

Pixel accuracy hardly changes, and is even slightly lower with the weights, because the gaps are under 1% of the pixels; the instance count is what changes. Over seeds 0 to 9 the plain loss recovered 63.7–90.3% of the cells and the weighted loss 95.9–98.1%; seed 8 is typical of both. With PyTorch’s default initialization in place of the Gaussian one, one weighted run in ten (seed 3) collapsed to predicting background everywhere.

Inference

The network is applied tile by tile as above and outputs a softmax map per class. For the EM challenge the authors averaged the predictions for seven rotated versions of the input, without further pre- or post-processing. Where instances are needed, the paper does not spell out the step; the natural one (our reading) is to take connected components of the thresholded foreground, which the gaps keep apart, as in the example above. They reported that a $512 \times 512$ image took less than a second on a recent GPU (2015).

Performance

Results from the paper (arXiv version; EM ranking of 6 March 2015):

Benchmark	Metric	U-Net	Comparison
ISBI 2012 EM segmentation (neuronal membranes, 30 training images)	Warping error	0.000353	0.000420 (Ciresan et al., IDSIA)
ISBI 2012 EM segmentation	Rand error	0.0382	0.0504 (IDSIA)
ISBI cell tracking challenge 2015, PhC-U373 (phase contrast, 35 partially annotated images)	IoU	0.9203	0.83 (second-best 2015)
ISBI cell tracking challenge 2015, DIC-HeLa (DIC, 20 partially annotated images)	IoU	0.7756	0.46 (second-best 2015)

On the EM challenge U-Net had the best warping error at the time; a few entries had lower Rand error, which the authors attribute to dataset-specific post-processing applied to Ciresan et al.’s probability maps. The text gives the warping error as 0.0003529 and Table 1 as 0.000353; for DIC-HeLa, the text gives an IoU of 77.5% and Table 2 0.7756, which rounds to 77.6%. The authors won both cell-tracking categories.

Strengths

• Learns from tens of annotated images, given elastic augmentation.
• Concatenated skips give boundaries at full resolution with little extra machinery.
• Fully convolutional: any valid input size, and arbitrarily large images by tiling.
• Simple enough to be adapted to 3D, to other losses and to other tasks.

Limitations

• Valid convolutions constrain tile sizes and require cropping, tiling and mirroring; many modern versions pad instead.
• The receptive field is bounded by the depth; long-range context needs more levels or attention.
• Separating touching objects relies on gaps in the labels and a hand-designed weight map; it fails where no gap pixel is visible, and dedicated instance methods such as Mask R-CNN or watershed post-processing take over.
• A fixed architecture still needs dataset-specific choices of patch size, spacing and normalization; nnU-Net was built to automate them.

Historical Importance

U-Net appeared in 2015 alongside FCN, and the two set the template for convolutional segmentation. U-Net’s contributions were the symmetric decoder, the concatenating skip connection, and a demonstration that a deep network could be trained from very few labeled images. It became a standard design in biomedical segmentation, and its encoder–decoder with skips spread far beyond: pix2pix uses a U-Net generator for image-to-image translation (encoder–decoder), and denoising diffusion probabilistic models represent their reverse, denoising process with a U-Net backbone (Ho, Jain & Abbeel, 2020).

Variants

• 3D U-Net (Çiçek et al., 2016, MICCAI) replaces every 2D operation with its 3D counterpart, adds batch normalization before each ReLU, and learns from sparsely annotated volumes by giving unlabeled voxels zero weight in the loss. Its input tile is $132 \times 132 \times 116$ voxels, its output $44 \times 44 \times 28$.
• V-Net (Milletari, Navab & Ahmadi, 2016) is a volumetric relative with residual stages, $5 \times 5 \times 5$ convolutions, and strided convolutions instead of pooling; it proposed a loss based on the Dice coefficient.
• Attention U-Net (Oktay et al., 2018, presented at MIDL 2018) multiplies the skip features by attention coefficients computed from the skip features and a coarser gating signal, suppressing irrelevant regions before concatenation; it was evaluated against a 3D U-Net on two abdominal CT datasets, one annotated for pancreas, liver and spleen and one for the pancreas alone.
• UNet++ (Zhou et al., 2018, DLMIA workshop) replaces each plain skip by nested, densely connected convolution blocks meant to narrow the semantic gap between encoder and decoder features, trains with deep supervision, and can be pruned at inference. Across four medical tasks the authors reported that UNet++ with deep supervision gained 3.9 IoU points on average over U-Net and 3.4 over a wider U-Net.
• nnU-Net (Isensee et al., 2021, Nature Methods) keeps the U-Net and automates everything around it: preprocessing, network topology, training and post-processing are configured from the dataset’s properties by fixed parameters, interdependent rules and empirical decisions. Without manual intervention it surpassed most existing approaches, including specialized ones, on 23 public datasets from international segmentation competitions. Its arXiv preprint (v2) describes its default 2D, 3D and cascaded 3D U-Nets with instance normalization, leaky ReLU, strided-convolution downsampling, transposed-convolution upsampling, deep supervision and a sum of cross-entropy and Dice losses.

Common modifications in today’s U-Nets (general practice rather than any one paper): padded convolutions so the output matches the input, batch or instance normalization, transposed convolution or bilinear upsampling followed by convolution, residual blocks within each stage, and pretrained classification backbones as the encoder.

Implementations

The authors released a Caffe implementation and trained networks with the paper, and nnU-Net is distributed as an open-source tool. U-Nets are short to write in any framework; the example above shows the parts: double-convolution blocks, pooling, up-convolution, concatenation and a $1 \times 1$ head. When following the original, keep the unpadded convolutions with the cropping and tile-size constraints; many reimplementations pad, so the output has the input’s size and tiling needs no 184-pixel margin, although, in our reading, zero-padded border pixels then see less real context than the mirrored tiles provide.