Lesson 41: Semantic Segmentation¶

Classification labels a whole image. Detection (Lessons 39-40) labels a handful of boxes. Semantic segmentation goes one step further: label every pixel with a class. This lesson builds a small U-Net-style encoder-decoder (Ronneberger, Fischer & Brox, 2015★) from scratch, and shows concretely why the architecture's defining feature — skip connections between matching encoder and decoder resolutions — matters, not just as a Lesson-35-style gradient-flow trick but for preserving spatial detail that pooling destroys.

In [1]:

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

A 3-class pixel labeling task¶

Each scene has several small circles and squares scattered on a background. The target isn't one label per image — it's a full per-pixel class map: 0 = background, 1 = circle, 2 = square.

In [2]:

def make_scene(rng, size=32, n_shapes=5):
    scene = np.zeros((size, size), dtype=np.float32)
    mask = np.zeros((size, size), dtype=np.int64)  # 0=background, 1=circle, 2=square
    for _ in range(n_shapes):
        shape_type = rng.choice([1, 2])
        r = rng.integers(2, 4)
        cx, cy = rng.integers(r, size - r), rng.integers(r, size - r)
        yy, xx = np.mgrid[0:size, 0:size]
        if shape_type == 1:
            m = ((xx - cx) ** 2 + (yy - cy) ** 2) <= r ** 2
        else:
            m = (np.abs(xx - cx) <= r) & (np.abs(yy - cy) <= r)
        scene[m] = 1.0
        mask[m] = shape_type
    scene = np.clip(scene + rng.normal(0, 0.05, scene.shape), 0, 1).astype(np.float32)
    return scene, mask

rng = np.random.default_rng(13)
N = 300
scenes, masks = [], []
for _ in range(N):
    s, m = make_scene(rng)
    scenes.append(s); masks.append(m)
scenes = np.array(scenes, dtype=np.float32)
masks = np.array(masks, dtype=np.int64)

split = int(0.85 * N)
Xtr, Mtr = scenes[:split], masks[:split]
Xte, Mte = scenes[split:], masks[split:]

print('class pixel fractions (train):', dict(zip(['background', 'circle', 'square'],
      (np.bincount(Mtr.ravel()) / Mtr.size).round(3))))

fig, axes = plt.subplots(2, 4, figsize=(9, 4.5))
for i in range(4):
    axes[0, i].imshow(Xtr[i], cmap='gray'); axes[0, i].axis('off')
    axes[1, i].imshow(Mtr[i], cmap='viridis', vmin=0, vmax=2); axes[1, i].axis('off')
axes[0, 0].set_title('input', fontsize=9, loc='left')
axes[1, 0].set_title('per-pixel label', fontsize=9, loc='left')
plt.show()

class pixel fractions (train): {'background': np.float64(0.866), 'circle': np.float64(0.046), 'square': np.float64(0.088)}

No description has been provided for this image

An encoder-decoder with skip connections¶

A segmentation network needs an output the same spatial size as the input, but with a class-probability vector at every pixel instead of one RGB value. The encoder half is an ordinary CNN, downsampling twice via max pooling (Lesson 33) to build up wide-receptive-field, semantically rich features (Lesson 11's pyramid, again). The decoder half upsamples back to full resolution. A plain encoder-decoder would only have the coarse, heavily-pooled bottleneck features to work with when reconstructing full resolution — so at each decoder stage, this network also concatenates in the encoder's feature map from the matching resolution, before pooling destroyed that detail. This is the "skip connection" that gives U-Net its name (and its U-shaped diagram) — architecturally similar to Lesson 35's residual connections, but concatenating features across the encoder/decoder divide rather than adding a delta within a single stack.

In [3]:

class UNetTiny(nn.Module):
    def __init__(self, n_classes=3):
        super().__init__()
        self.enc1 = nn.Sequential(nn.Conv2d(1, 16, 3, padding=1), nn.ReLU())
        self.enc2 = nn.Sequential(nn.Conv2d(16, 32, 3, padding=1), nn.ReLU())
        self.enc3 = nn.Sequential(nn.Conv2d(32, 64, 3, padding=1), nn.ReLU())
        self.pool = nn.MaxPool2d(2)
        self.up = nn.Upsample(scale_factor=2, mode='nearest')
        self.dec2 = nn.Sequential(nn.Conv2d(64 + 32, 32, 3, padding=1), nn.ReLU())
        self.dec1 = nn.Sequential(nn.Conv2d(32 + 16, 16, 3, padding=1), nn.ReLU())
        self.out = nn.Conv2d(16, n_classes, 1)

    def forward(self, x):
        f1 = self.enc1(x)                                     # (B,16,H,W)
        f2 = self.enc2(self.pool(f1))                          # (B,32,H/2,W/2)
        f3 = self.enc3(self.pool(f2))                          # (B,64,H/4,W/4)
        d2 = self.dec2(torch.cat([self.up(f3), f2], dim=1))    # skip from f2
        d1 = self.dec1(torch.cat([self.up(d2), f1], dim=1))    # skip from f1
        return self.out(d1)

torch.manual_seed(0)
model = UNetTiny()
opt = torch.optim.Adam(model.parameters(), lr=0.01)
Xt = torch.tensor(Xtr).unsqueeze(1); Mt = torch.tensor(Mtr)
for _ in range(200):
    opt.zero_grad()
    loss = F.cross_entropy(model(Xt), Mt)
    loss.backward()
    opt.step()

with torch.no_grad():
    preds = model(torch.tensor(Xte).unsqueeze(1)).argmax(1).numpy()

def mean_iou(preds, targets, n_classes=3):
    ious = []
    for c in range(n_classes):
        p, t = preds == c, targets == c
        inter, union = (p & t).sum(), (p | t).sum()
        ious.append(inter / union if union > 0 else float('nan'))
    return ious

pixel_acc = (preds == Mte).mean()
ious = mean_iou(preds, Mte)
print(f'pixel accuracy: {pixel_acc:.1%}')
for name, iou in zip(['background', 'circle', 'square'], ious):
    print(f'  {name:>10} IoU: {iou:.3f}')
print(f'mean IoU: {np.mean(ious):.3f}')

pixel accuracy: 99.5%
  background IoU: 1.000
      circle IoU: 0.893
      square IoU: 0.937
mean IoU: 0.943

Do the skip connections actually matter?¶

Train an otherwise-identical network with the skip connections removed — the decoder only ever sees the pooled, upsampled bottleneck features, never the original-resolution encoder features.

In [4]:

class UNetNoSkip(nn.Module):
    def __init__(self, n_classes=3):
        super().__init__()
        self.enc1 = nn.Sequential(nn.Conv2d(1, 16, 3, padding=1), nn.ReLU())
        self.enc2 = nn.Sequential(nn.Conv2d(16, 32, 3, padding=1), nn.ReLU())
        self.enc3 = nn.Sequential(nn.Conv2d(32, 64, 3, padding=1), nn.ReLU())
        self.pool = nn.MaxPool2d(2)
        self.up = nn.Upsample(scale_factor=2, mode='nearest')
        self.dec2 = nn.Sequential(nn.Conv2d(64, 32, 3, padding=1), nn.ReLU())
        self.dec1 = nn.Sequential(nn.Conv2d(32, 16, 3, padding=1), nn.ReLU())
        self.out = nn.Conv2d(16, n_classes, 1)

    def forward(self, x):
        f1 = self.enc1(x)
        f2 = self.enc2(self.pool(f1))
        f3 = self.enc3(self.pool(f2))
        d2 = self.dec2(self.up(f3))    # no skip: only the pooled, upsampled features
        d1 = self.dec1(self.up(d2))    # no skip
        return self.out(d1)

torch.manual_seed(0)
model_noskip = UNetNoSkip()
opt2 = torch.optim.Adam(model_noskip.parameters(), lr=0.01)
for _ in range(200):
    opt2.zero_grad()
    loss = F.cross_entropy(model_noskip(Xt), Mt)
    loss.backward()
    opt2.step()

with torch.no_grad():
    preds_noskip = model_noskip(torch.tensor(Xte).unsqueeze(1)).argmax(1).numpy()

pixel_acc_noskip = (preds_noskip == Mte).mean()
ious_noskip = mean_iou(preds_noskip, Mte)

print(f'{"":>18} {"pixel acc":>10} {"mean IoU":>10}')
print(f'{"no skip":>18} {pixel_acc_noskip:>10.1%} {np.mean(ious_noskip):>10.3f}')
print(f'{"with skip":>18} {pixel_acc:>10.1%} {np.mean(ious):>10.3f}')

                    pixel acc   mean IoU
           no skip      98.1%      0.870
         with skip      99.5%      0.943

In [5]:

fig, axes = plt.subplots(4, 4, figsize=(9, 9))
for i in range(4):
    axes[0, i].imshow(Xte[i], cmap='gray')
    axes[1, i].imshow(Mte[i], cmap='viridis', vmin=0, vmax=2)
    axes[2, i].imshow(preds[i], cmap='viridis', vmin=0, vmax=2)
    axes[3, i].imshow(preds_noskip[i], cmap='viridis', vmin=0, vmax=2)
    for r in range(4):
        axes[r, i].axis('off')
for r, name in enumerate(['input', 'true mask', 'pred (skip)', 'pred (no skip)']):
    axes[r, 0].set_title(name, fontsize=9, loc='left')
plt.tight_layout()
plt.show()

With these small, closely-packed shapes, removing the skip connections costs several points of mean IoU. The reason is exactly what the architecture predicts: after two rounds of pooling, the bottleneck has only a quarter of the spatial resolution, and small shapes (radius 2-3 pixels) can blur together or lose their boundaries entirely at that resolution. Upsampling a blurry, low-resolution guess doesn't recover the lost detail — nearest-neighbor upsampling (or any fixed interpolation) can only spread existing information around, not invent missing edges. The skip connection sidesteps the problem by handing the decoder the original-resolution features directly, so precise boundaries never had to survive the bottleneck in the first place.

Exercise¶

Increase n_shapes from 5 to 10, making the scene more crowded. Does the skip-vs-no-skip gap in mean IoU get larger or smaller? What does that suggest about when skip connections matter most?
This lesson's loss is plain per-pixel cross-entropy. Print the per-pixel class weights implied by class pixel fractions above and try F.cross_entropy(logits, Mt, weight=inverse_class_freq) (Lesson 37's imbalance fix, applied here) to see whether it changes the circle/square IoU balance.
nn.Upsample(mode='nearest') was used for simplicity. Try mode='bilinear', align_corners=False instead (Lesson 9's bilinear interpolation, now inside a network) and compare mean IoU for both the skip and no-skip models. Does the smoother upsampling help more or less than the skip connection does?

← Lesson 40: Object Detection II — From Classification to Regression

Lesson 42: Instance Segmentation →