U-Net: Convolutional Networks for Biomedical Image Segmentation

A convolutional network that labels every pixel of a microscope image, trained on a few dozen annotated examples.

A contracting path shrinks the image to gather context, an expanding path grows it back to full resolution, and skip connections copy each level's high-resolution features across so that boundaries land on the right pixels.

In 2015 the strongest method for segmenting electron-microscope images ran a classifier once per pixel. U-Net labels a 388×388 tile in one forward pass, and it trained on 20 to 35 annotated images per task.

An image classifier maps a whole picture to one label. By 2015 convolutional networks did this well, after training on ImageNet's million labelled images. Biomedical segmentation asks for a class at every pixel instead. In an electron-microscope slice of nerve tissue each pixel is either the inside of a cell or a membrane, and where two cells press together the membrane between them is a line one or two pixels wide. A segmentation that misses such a line merges two cells into one.

Deciding whether a dark line is a membrane or a fold in the tissue section takes context: the answer shows only over a wide area around the pixel. A convolutional network collects context by pooling. A $2\times2$ max-pool replaces each 2×2 block of a feature map with its largest value, so after four of them, with the convolutions in between, a single position in the map depends on a 140×140 patch of the input. The same four pools have divided the map's width by 16, and the network has discarded the record of which pixel inside that patch produced a feature. More context costs localization, the ability to say which pixel a boundary is on.

U-Net, from Olaf Ronneberger, Philipp Fischer and Thomas Brox at the University of Freiburg, answers with three architectural pieces and three training choices:

• a contracting path that pools, as a classifier does, to gather context (the left arm of the U);
• an expanding path that upsamples back toward full resolution (the right arm);
• skip connections that copy each contracting level's feature map across to the expanding level of the same size;
• unpadded convolutions and an overlap-tile scheme, so images of any size are segmented from real context only;
• a weighted loss that adds up to 10 to the weight of pixels in the thin borders between touching cells;
• elastic deformations that turn 30 training images into an unlimited supply of warped variants.

The sliding-window baseline

The method U-Net set out to replace is the sliding-window network of Ciresan and colleagues, which won the ISBI 2012 electron-microscopy (EM) segmentation challenge by a large margin. To label pixel $(i, j)$, it cuts out a square patch centred on that pixel and asks a classifier network whether the centre is membrane. Repeating that for every pixel gives a full segmentation. The approach has one strength the paper points out: every pixel is a training example, so 30 annotated images become millions of labelled patches.

It also has two drawbacks, both named in the paper. The first is speed. A 512×512 image has 262,144 pixels, so labelling it takes 262,144 forward passes, and two neighbouring patches share all but one column of their pixels, so almost all of that computation is repeated. The second is a trade-off set by the patch size. A larger patch shows more context but needs more pooling layers to reduce it to one prediction, and the pooling blurs where the centre pixel's boundary lies. A smaller patch localizes well but shows too little of the surroundings. Drag the patch size in Figure 1 and watch the two curves cross:

The curves in Figure 1 are schematic; the next section computes the real quantities, how far each layer's view reaches and how coarse its map becomes. U-Net keeps one idea from the fully convolutional network (FCN) of Long, Shelhamer and Darrell: no classifier runs per pixel. Every layer is a convolution, so the network accepts a whole tile and produces a whole label map in one pass, and the computation shared by neighbouring pixels is done once.

The contracting path: context from pooling

The left arm of the U is an ordinary convolutional network. Each of its levels runs two $3\times3$ convolutions, each followed by a ReLU (which sets negative values to zero), then a $2\times2$ max-pool with stride 2. Every pool halves the height and width, and the number of feature channels, the values stored at each pixel position, doubles at each step down: 64, 128, 256, 512, and 1024 at the bottom level, which is called the bottleneck. The convolutions are unpadded, meaning they add no zeros at the border, so each one returns a map 2 pixels narrower than its input. In the paper's example a 572×572 input tile becomes 568, 280, 136, 64 and finally 28 pixels wide at the bottleneck.

The receptive field of a neuron is the patch of the input tile its value depends on, and it grows by a fixed rule. A $3\times3$ convolution widens it by 2 steps of the current map, and a step of the current map is 1 input pixel at the top level, 2 after one pool, 4 after two, and so on. So the two convolutions at the top give a 5-pixel receptive field. After the first pool each convolution adds 4 input pixels (the pool itself adds one step of the map it reads), and the field reaches 14 at level 1, then 32 at level 2, 68 at level 3 and 140 at the bottleneck, about a quarter of the 572-pixel tile. Over the same levels the map shrinks from 568 to 28 pixels wide. Step the depth slider in Figure 2 from the top level to the bottleneck:

At the bottleneck a neuron's value depends on 140×140 input pixels. Neighbouring bottleneck positions sit 16 input pixels apart, though, so a bottleneck feature cannot distinguish a boundary on one pixel from a boundary on the next.

The expanding path and the full U

The right arm reverses the contracting path, one level at a time. Each step does four things, in this order:

1. an up-convolution: upsample the map by 2 in each direction, then apply a $2\times2$ convolution that halves the number of channels (1024 to 512 at the first step);
2. take the feature map saved at the same level of the contracting path and crop out its centre to the size of the upsampled map;
3. concatenate the two along the channel axis, which doubles the channels again;
4. two unpadded $3\times3$ convolutions with ReLU.

Take the first step up. The 28×28×1024 bottleneck upsamples to 56×56 with 512 channels. The saved level-3 map is 64×64×512, so its centre 56×56 is cropped out, 4 pixels from each side, and stacked onto the upsampled map to make 56×56×1024. Two convolutions bring it to 52×52×512. Three more steps like it climb through 100, 196 and 388 pixels. At the top a $1\times1$ convolution turns each pixel's 64 features into one score per class, 2 classes in the paper's example (for EM, cell and membrane). Counting layers: 18 of the 3×3 convolutions (8 down, 2 at the bottleneck, 8 up), 4 up-convolutions and the final 1×1 make the 23 convolutional layers the paper reports. Press Play in Figure 3, or scrub, to walk the eleven stages with their exact sizes:

In pseudocode the whole forward pass is two loops. The contracting loop appends each level's last map to a list before pooling, and the expanding loop pops them in reverse order:

# U-Net forward pass for one 572x572 tile; shapes are (channels, H, W)
skips = []
x = tile                                  # (1, 572, 572), mirror-padded
for ch in [64, 128, 256, 512]:            # contracting path
    x = relu(conv3x3(x, ch))              # valid conv: H, W shrink by 2
    x = relu(conv3x3(x, ch))              # (64, 568, 568) at the first level
    skips.append(x)                       # saved for the expanding path
    x = maxpool2x2(x)                     # (64, 284, 284): H, W halve
x = relu(conv3x3(x, 1024))                # bottleneck
x = relu(conv3x3(x, 1024))                # (1024, 28, 28)
for skip in reversed(skips):              # expanding path
    x = upconv2x2(x)                      # 2x upsample + 2x2 conv, channels / 2
    x = concat([crop_center(skip, x.H), x])   # channels double again
    x = relu(conv3x3(x, x.C // 2))
    x = relu(conv3x3(x, x.C))             # (64, 388, 388) after the last level
return conv1x1(x, n_classes)              # (2, 388, 388): per-pixel scores

The paper's up-convolution is "an upsampling of the feature map followed by a $2\times2$ convolution". Many later implementations replace it with one strided transposed convolution, which also doubles the resolution with learned weights. The two behave alike but are not identical: in a strided transposed convolution neighbouring output pixels receive contributions from different numbers of kernel taps, which can print a faint checkerboard into the output, and resize-then-convolve does not have that pattern (Odena, Dumoulin and Olah, 2016).

The U shape comes from one change to FCN that the paper names. FCN upsamples class-score maps, one channel per class. U-Net keeps a large number of feature channels in the expanding path, 512 down to 64, so context from the deep layers can flow up to the high-resolution ones. With that many channels on the way up, the expanding path mirrors the contracting path.

What the skip connections carry

Without the skips, the expanding path would have only the bottleneck to work from. A 28×28 map with positions 16 input pixels apart locates a cell only to within 16 pixels, too coarse to place a membrane 2 pixels wide or to separate two touching cells when the gap between them is narrower than its grid. Upsampling adds pixels without adding information: the upsampled map is a smooth interpolation of the coarse one.

The information is still in the network, in the contracting path's early maps. The level-0 map is 568×568, computed before any pooling, and records edges at single-pixel precision. The skip connection sends that saved map across to the expanding path, where it arrives at the matching resolution. Each level of skip restores a factor of 2 in resolution: with no skips the output boundary can at best be placed on a 16-pixel grid, with the level-3 skip on an 8-pixel grid, and with all four on the 1-pixel grid of the input.

Figure 4 models that argument directly. It takes a true mask of two touching cells, separated by a background border about 2 pixels wide, averages it onto the grid of the finest level that reaches the output, upsamples it smoothly and thresholds it. Start with no skips, then connect them one level at a time from the bottom up, and watch the cell count:

With no skips, the model loses 13% of the cell area to errors and finds one cell where there are two. With all four the pixel error is 0 and both cells are found. The intermediate settings show that the separating border needs the finest skip: at a 2-pixel grid the pixel error is only 1.1%, but that 1.1% includes the border, and the two cells merge.

U-Net joins the skip to the upsampled map by concatenation, where FCN added its coarse and fine predictions together. After concatenation the next $3\times3$ convolution has separate weights for the 64 skip channels and the 64 upsampled channels, so it learns per output feature how much of each to use. After an addition, the next convolution has one weight per summed channel and must apply it to the skip value and the upsampled value alike. Concatenation doubles the input channels instead: the first convolution after the top-level concat reads 128 channels and has $3 \cdot 3 \cdot 128 \cdot 64$, or 73,728, weights, twice the 36,864 it would need after an addition.

The word "skip" also names the additive shortcut in ResNet, which is a different mechanism: a ResNet shortcut adds a block's input to its output at the same resolution to help optimization, and a U-Net skip carries a saved map across the U to restore spatial detail.

Valid convolutions and the overlap-tile strategy

Every convolution in U-Net keeps only its valid part, the positions where the whole $3\times3$ kernel lies on real input. Padding with zeros would keep the map size, but the pixels near the border would then be computed partly from zeros, a hard black frame that real tissue does not have. With valid convolutions every output pixel is computed from real image context only, and the price is a shrinking map: a 572×572 tile yields a 388×388 segmentation, 184 pixels narrower, which is 92 pixels of context on each side.

The shrinkage is also why the skips are cropped. The level-0 map is 568 pixels wide, but by the time the expanding path has come back up to level 0 it has lost border pixels at every level below, and the upsampled map is 392 wide. Only the centre 392×392 of the saved map lines up with it, so the crop removes 88 pixels from each side. The crops at the four levels are 88, 40, 16 and 4 pixels per side.

The pools add a constraint on the tile size. The paper says the input tile must be chosen so that every $2\times2$ max-pool is applied to a map with even width and height. If a pool meets an odd width, a framework drops the last column (or pads one), the upsampled map is then an odd number of pixels narrower than the saved map, and the centre crop would have to start at a half pixel. Output tiles computed that way do not line up with their neighbours. Working the rules through, the valid input sizes are 188, 204, 220 and so on in steps of 16; 572 is one of them, and every valid tile gives an output 184 pixels narrower. Drag the input size in Figure 5 through valid and invalid values:

Because the output is smaller than the input, a large image is segmented in tiles. Output tiles of 388×388 are laid edge to edge; each one needs a 572×572 input tile centred on it, so neighbouring input tiles overlap by 184 pixels. At this tile size, four output tiles cover a 512×512 EM image, which is four forward passes against the 262,144 of the sliding-window network. Tiling also makes the memory cost depend on the tile size alone; the paper notes that without it the resolution would be limited by GPU memory.

At the edge of the image, the input tile needs 92 pixels of context that do not exist. U-Net fills them by mirroring: the image is reflected across its own border, so the missing strip is a mirror copy of the pixels just inside. The reflected strip has the same texture and brightness statistics as the tissue, where a zero pad would add a black edge. Slide the tile in Figure 6 to either edge:

The weighted loss

The network's last layer produces $K$ scores $a_k(\mathbf{x})$ at each output pixel $\mathbf{x}$, one per class. A softmax over the classes, pixel by pixel, turns them into probabilities:

The largest score gets a probability near 1 and the others near 0. With $\ell(\mathbf{x})$ the true label at pixel $\mathbf{x}$ and $\Omega$ the set of output pixels, training minimizes a weighted cross-entropy, which the paper calls the energy function:

Each pixel contributes $-\log$ of the probability assigned to its correct class, times a weight. A pixel predicted correctly with probability 0.9 costs 0.105; at 0.5 it costs 0.693; at 0.1 it costs 2.303. The paper prints Eq (1) without the leading minus sign. Since a probability is at most 1, its log is at most 0, so the printed sum is never positive and would have to be maximized; the loss that training minimizes carries the minus, as written above and in the code below.

The weight $w(\mathbf{x})$ handles a problem specific to cell images: two touching cells of the same class. If every pixel had weight 1, a network that painted both cells and the 2-pixel border between them as foreground would pay for only those few border pixels, a small fraction of the image, and the two cells would come out as one. The paper therefore precomputes, for each ground-truth segmentation, a weight map that makes the border pixels expensive:

The first term, $w_c$, balances class frequencies, so that a class with few pixels is not outweighed by a common one. In the second term $d_1(\mathbf{x})$ is the distance from pixel $\mathbf{x}$ to the border of the nearest cell and $d_2(\mathbf{x})$ the distance to the border of the second-nearest cell. The paper sets $w_0$ to 10 and $\sigma$ to about 5 pixels.

Take a background pixel on the straight line between two round cells. Its distances to the two borders add up to exactly the gap $g$ between the cells, wherever it sits on that line. With $w_c$ at 1, the weight across the gap is $1 + 10\,e^{-g^2/50}$:

• a 2-pixel gap: 10.23, about ten times an ordinary background pixel;
• a 5-pixel gap ($g$ equal to $\sigma$): the added term is $10\,e^{-0.5}$, about 6.07, so the weight is about 7.07;
• a 10-pixel gap: 2.35;
• a 20-pixel gap: 1.003, no measurable bonus.

Background far from any pair of cells has a large $d_1 + d_2$ and keeps its class weight. The Gaussian makes the fall-off fast: the bonus is above half its peak only while $d_1 + d_2$ stays below about 1.18σ, about 6 pixels. In Figure 7, move the cells together and apart, change σ, and hover over the map to read $d_1$, $d_2$ and $w$ at any pixel:

The map depends only on the labels, so it is a constant during training and receives no gradient. A training step runs one tile through the network, takes the per-pixel cross-entropy against the labels, multiplies by the precomputed weights and sums:

# one training step (batch = one tile); w is precomputed per label map
scores = unet(tile)                       # (K, 388, 388)
logp = log_softmax(scores, dim=0)         # log p_k(x) at every pixel
nll = -logp.gather(0, label[None])[0]     # -log p_{l(x)}(x), shape (388, 388)
loss = (w * nll).sum()                    # Eq (1); w is a constant, no gradient
loss.backward()                           # gradients reach every conv weight
sgd_momentum_step(lr, momentum=0.99)

Training on 30 images

The EM training set is 30 images of 512×512 pixels, serial-section transmission electron micrographs of the ventral nerve cord of a Drosophila first-instar larva. A network with 23 convolutional layers can memorize 30 images. The paper answers with data augmentation, and names the invariances microscopy needs: shifts, rotations, deformations and changes in grey value. Of these it calls random elastic deformation "the key concept" for training with very few annotated images. Rotations and shifts move the image as a rigid whole; tissue in a section stretches and bends locally, and elastic warps produce that kind of variation.

The recipe has three steps. Place a coarse 3×3 grid of control points over the image and give each a random displacement drawn from a Gaussian with a standard deviation of 10 pixels. Compute the displacement at every pixel by bicubic interpolation between the control points, which makes the warp smooth. Apply the same warp to the image and to its label map, so they stay aligned. Each draw is a new training image with a correct label. Press for new samples in Figure 8 and move the standard deviation from 0 to twice the paper's value:

Elastic deformation of training images is older than U-Net; it goes back at least to Simard, Steinkraus and Platt's handwriting work in 2003. U-Net's own citation for the value of augmentation is Dosovitskiy et al. The paper adds one more source of variation: dropout layers at the end of the contracting path, which zero a random subset of activations during training and which it describes as further implicit data augmentation.

The remaining training settings follow from small data on a small GPU. Training uses the stochastic gradient descent of the Caffe framework with a batch of a single image: the authors prefer a large input tile, with more context per pass, over a large batch. To make up for the noise of one-image gradients they set the momentum to 0.99, "such that a large number of the previously seen training samples determine the update in the current optimization step". One way to read that number (this page's reading, not the paper's): momentum keeps a running sum in which a gradient from $k$ steps ago has weight $0.99^k$. Those weights sum to $1/(1-0.99)$, which is 100, and a gradient's weight halves after about 69 steps, so each update blends roughly the last hundred images' gradients.

Initialization matters in a network with 23 layers and several paths through it; the paper warns that otherwise some parts give excessive activations while others never contribute. The weights are drawn from a Gaussian with standard deviation $\sqrt{2/N}$, where $N$ is the number of inputs feeding one neuron, so that each feature map starts with roughly unit variance. For a 3×3 convolution reading 64 channels:

This is the rule of He et al. (2015). The factor 2 corrects for the ReLU, which sets about half of its inputs to zero and so halves the variance passed on; in He et al.'s analysis the linear-case derivation of Glorot and Bengio gives $\sqrt{1/N}$. Training took about 10 hours on an NVIDIA Titan GPU with 6 GB of memory.

Results on three datasets

The EM challenge scores a submitted membrane-probability map by thresholding it at 10 levels and computing three errors, all lower-is-better: warping error, which counts topological mistakes such as splits and merges of cells; Rand error, which measures disagreement over whether pairs of pixels belong to the same object; and pixel error, the plain fraction of mislabelled pixels. The paper's Table 1 is the leaderboard of March 6th, 2015, sorted by warping error.

U-Net, with its prediction averaged over 7 rotated versions of the input and no pre- or post-processing, reached a warping error of 0.0003529, the new best score. Ciresan et al.'s best sliding-window submission had 0.000420, and the runner-up, DIVE-SCI, 0.000355, so on the ranking metric U-Net leads second place by 0.000002. On Rand error, though, U-Net's 0.0382 is beaten by DIVE-SCI (0.0305) and IDSIA-SCI (0.0189); the paper says those better Rand scores come from highly dataset-specific post-processing applied to Ciresan et al.'s probability map, and a footnote adds that the authors of one of them submitted 78 different solutions. On pixel error U-Net's 0.0611 is also above DIVE-SCI's 0.0584 and DIVE's 0.0582.

The second benchmark is the ISBI cell tracking challenge 2015, scored by intersection over union (IoU), the overlap between the predicted and the true mask divided by their union, so higher is better. On PhC-U373, glioblastoma cells under phase-contrast microscopy with 35 partially annotated training images, U-Net reached an IoU of 92% against 83% for the second-best 2015 entry. On DIC-HeLa, HeLa cells under differential interference contrast with 20 partially annotated training images, it reached 77.5% against 46% (Table 2 gives 0.7756). It won both categories by a large margin. Switch between the three views in Figure 9:

So on EM, U-Net ranks first by the metric the leaderboard is sorted on and beats the sliding-window network on all three errors, without post-processing; on the cell-tracking datasets it wins outright, and the DIC-HeLa margin is 31.5 points of IoU. Inference is fast: a 512×512 image takes less than a second on a recent GPU, against one forward pass per pixel for the method it replaced.

The same architecture, contracting path plus expanding path plus concatenated skips, became the default for dense prediction, where an output is needed at every pixel. The noise-prediction network in DDPM is a U-Net, and so is the denoiser in Latent Diffusion, which runs it on a compressed latent image. Those versions use same-padding, normalization layers such as batch normalization or its group-wise variant, and attention blocks, but they still save each level's map on the way down and concatenate it back in on the way up.