Segmentation Architectures: U-Net

Build the mental model

Classification networks intentionally throw away spatial resolution as they go deeper — pooling and strided convolutions shrink the feature map down to a small, highly abstract representation, then a final layer collapses it to a single label per image. Segmentation needs the opposite output: a label for every pixel, at the original image resolution. A naive encoder-only network can't do this well, because by the time the features are abstract enough to know 'this region is a cat', all the fine spatial detail about exactly where the cat's edges are has already been discarded by pooling. U-Net's architecture is a symmetric encoder-decoder: the encoder (contracting path) downsamples as usual to build up semantic understanding, and a mirrored decoder (expansive path) upsamples step by step, using transposed convolutions or interpolation, back to the original resolution.

The part that makes U-Net actually work well, rather than just producing a blurry upsampled blob, is the skip connections between mirrored encoder and decoder stages: at each decoder step, the upsampled feature map is concatenated (not added, like ResNet) with the encoder feature map from the same spatial resolution, before being passed through more convolutions. This gives the decoder direct access to the high-resolution spatial detail that was captured early in the encoder — sharp edges and fine textures — while still benefiting from the coarse, semantically rich features that came from the bottleneck. Without these skip connections, the decoder would have to somehow reconstruct precise pixel boundaries purely from the heavily compressed bottleneck representation, which loses far too much spatial information to do so accurately.

Connect it to a real scenario

The Tutorial Platform could use a U-Net-style segmentation model to produce a pixel-precise mask over sensitive regions in user-submitted code screenshots — API keys, tokens, or personal file paths pasted into a terminal — since only a full-resolution, pixel-level output (not a coarse bounding box) lets the platform blur exactly the offending text before publishing the screenshot without obscuring the surrounding code.

Classification networks discard spatial resolution as they deepen, which makes a plain encoder-only network unsuitable for pixel-level segmentation output.
U-Net's decoder mirrors the encoder, upsampling step by step back to the original spatial resolution instead of collapsing to a single label.
Skip connections concatenate (not add) matching-resolution encoder features into the decoder, restoring fine spatial detail lost during downsampling.

Try the working example

python

import torch
import torch.nn as nn
import torch.nn.functional as F

class MiniUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.enc1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.pool1 = nn.MaxPool2d(2)
        self.enc2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.pool2 = nn.MaxPool2d(2)

        self.bottleneck = nn.Conv2d(32, 64, kernel_size=3, padding=1)

        self.up2 = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
        self.dec2 = nn.Conv2d(64, 32, kernel_size=3, padding=1)   # 32 (up2) + 32 (skip) = 64 in
        self.up1 = nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2)
        self.dec1 = nn.Conv2d(32, 16, kernel_size=3, padding=1)   # 16 (up1) + 16 (skip) = 32 in

        self.out_conv = nn.Conv2d(16, 1, kernel_size=1)

    def forward(self, x):
        e1 = F.relu(self.enc1(x))          # (batch, 16, H, W)
        p1 = self.pool1(e1)                # (batch, 16, H/2, W/2)
        e2 = F.relu(self.enc2(p1))         # (batch, 32, H/2, W/2)
        p2 = self.pool2(e2)                # (batch, 32, H/4, W/4)

        b = F.relu(self.bottleneck(p2))    # (batch, 64, H/4, W/4)

        u2 = self.up2(b)                                       # (batch, 32, H/2, W/2)
        d2 = F.relu(self.dec2(torch.cat([u2, e2], dim=1)))      # skip connection with e2

        u1 = self.up1(d2)                                      # (batch, 16, H, W)
        d1 = F.relu(self.dec1(torch.cat([u1, e1], dim=1)))      # skip connection with e1

        return self.out_conv(d1)           # (batch, 1, H, W)

torch.manual_seed(0)
model = MiniUNet()
images = torch.randn(2, 3, 64, 64)  # fake batch of 2 RGB images
output = model(images)
print(output.shape)
print(output.shape[-2:] == images.shape[-2:])

You should see

torch.Size([2, 1, 64, 64])
True — after two downsampling stages (64x64 -> 32x32 -> 16x16) and two matching upsampling stages (16x16 -> 32x32 -> 64x64), the output spatial dimensions exactly match the 64x64 input, and the final print confirms this equality; the single output channel represents a per-pixel prediction map, the defining property that makes U-Net suitable for segmentation.

5-minute try-it

Add a third encoder/decoder level (enc3/pool3 downsampling to H/8, and a matching up3/dec3 with its own skip connection to e3) and confirm the final printed output shape is still (2, 1, 64, 64) — this shows the skip-connection pattern scales to arbitrarily deep U-Nets as long as every downsampling stage has a mirrored upsampling stage.

One important caution

Using an input spatial size that isn't evenly divisible by 4 (the total downsampling factor from two pool layers) makes the encoder and decoder feature maps end up with mismatched sizes, so torch.cat raises a dimension-mismatch error at the skip connection.

Swapping torch.cat for simple addition in the skip connection (mimicking ResNet) silently runs without error but discards information, since it forces the encoder and decoder channel counts to be equal instead of preserving both feature sets side by side.

Wikipedia — U-Net — Computer Vision