Lesson 110: Semantic Segmentation & U-Net

1. Semantic Segmentation & U-Net from Scratch

Title

USAAIO · Lesson 110 · Phase 3

Assign a class label to every pixel. Build the encoder-decoder with skip connections, three upsampling methods, and two imbalance-aware losses — all in PyTorch, all verified.

2. By the end of this lesson you can

Objectives

1. Define semantic segmentation and identify the output tensor shape (B, C, H, W)
2. Trace the U-Net encoder-decoder path and explain why skip connections use concatenation, not addition
3. Implement three upsampling methods — bilinear interpolation, ConvTranspose2d, pixel shuffle — and state the output-size formula for each
4. Implement Dice loss and Focal loss from scratch and explain when each beats pixel-wise cross-entropy
5. Build TinyUNet end-to-end and verify input/output shapes match (B, n_classes, H, W)

3. What survived from Convex Optimization?

Warm-up

4. Semantic segmentation — the task

Section

Part 1 of 4

5. Every pixel gets a label

Concept

Semantic segmentation is dense classification: instead of one label per image (Lessons 55-92), every pixel is assigned a class. The model output is a (B, C, H, W) tensor of per-class logits — one spatial map per class.

\[ \hat{y}_{b,c,i,j} = \text{logit for class } c \text{ at pixel } (i,j) \text{ in image } b \]

6. Fill in: input for Every pixel gets a label

Comparison

7. Encoder-decoder: the core idea

Concept

A plain CNN loses spatial resolution via pooling — useful for classification, fatal for segmentation. The encoder-decoder architecture restores it: downsample for context, then upsample to full resolution.

• Encoder: stack of Conv + MaxPool layers — shrinks H×W, grows channels (context)
• Bottleneck: deepest layer — maximum receptive field, smallest spatial size
• Decoder: upsample back up to (H, W) — recovers spatial resolution for per-pixel prediction

The problem: the decoder only sees a compressed bottleneck, so fine-grained edges and textures are lost. U-Net's solution is skip connections from each encoder level directly to the matching decoder level.

8. U-Net — skip connections & shapes

Section

Part 2 of 4

9. U-Net skip connections: concat, not add

Concept

U-Net (Ronneberger 2015) connects each encoder stage to the same-resolution decoder stage with a concatenation skip connection — the encoder feature map is appended along the channel dimension.

\[ \text{dec}_{in} = \text{concat}(\text{upsample}(\text{dec}_{prev}),\; \text{enc}_{same}) \quad \text{(channel dim)} \]

This doubles the channel count going into the decoder conv block, so the network can independently learn what from the encoder and what from the decoder are useful — the ResNet addition (Lesson 70) would force them to the same representation.

10. Guess the shape of the answer: U-Net tensor shapes — full trace

Estimation

11. U-Net tensor shapes — full trace

Worked example

Trace a U-Net with base=16 channels on a (B=2, 1, 32, 32) input. Predict each shape before reading on.

import torch, torch.nn as nn

class DoubleConv(nn.Module):
    def __init__(self, in_c, out_c):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_c, out_c, 3, padding=1),
            nn.BatchNorm2d(out_c), nn.ReLU(inplace=True),
            nn.Conv2d(out_c, out_c, 3, padding=1),
            nn.BatchNorm2d(out_c), nn.ReLU(inplace=True),
        )
    def forward(self, x): return self.conv(x)

torch.manual_seed(42)
x = torch.randn(2, 1, 32, 32)
enc1 = DoubleConv(1, 16);   e1 = enc1(x)           # no pool
enc2 = DoubleConv(16, 32);  e2 = enc2(nn.MaxPool2d(2)(e1))
bottle = DoubleConv(32, 64); b = bottle(nn.MaxPool2d(2)(e2))
print('e1:', e1.shape, '  e2:', e2.shape, '  bottle:', b.shape)

12. What each one costs: U-Net tensor shapes — full trace

Trade off

13. Something is wrong here: using addition instead of concatenation for U-Net skips

Anomaly

14. Trap: using addition instead of concatenation for U-Net skips

Trap

15. Break it on purpose: using addition instead of concatenation for…

Break the constraint

16. Upsampling methods — three ways to grow H×W

Section

Part 3 of 4

17. Three upsampling methods compared

Concept

The decoder must recover spatial resolution lost during encoding. Three methods are standard; they differ in whether upsampling weights are fixed or learned.

U-Net canonically uses ConvTranspose2d. Bilinear + Conv is also popular (avoids checkerboard artifacts). Pixel shuffle is common in super-resolution (Lesson 98 territory).

18. Break it if you can: Three upsampling methods compared

Counterexample

19. ConvTranspose2d — output size formula

Worked example

Predict the output size of ConvTranspose2d(4, 2, kernel_size=2, stride=2) applied to inputs of H=4, 8, 16. Use the formula before running.

\[ H_{out} = (H_{in} - 1) \cdot \text{stride} - 2 \cdot \text{padding} + \text{kernel\_size} \]

import torch, torch.nn as nn

ct = nn.ConvTranspose2d(4, 2, kernel_size=2, stride=2)
for h in [4, 8, 16]:
    x = torch.randn(1, 4, h, h)
    with torch.no_grad():
        out = ct(x)
    formula = (h - 1) * 2 - 0 + 2      # padding=0
    print(f'in={h:>2} -> out={out.shape[-1]:>2}  (formula {formula})')

20. All three upsamplers — verified shapes

Worked example

Apply all three methods to the same (2, 8, 4, 4) tensor and confirm shapes. Pixel shuffle requires a preceding conv to expand channels to C * r².

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

torch.manual_seed(42)
x = torch.randn(2, 8, 4, 4)

# Bilinear (no params)
up_bil = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
print('bilinear:     ', up_bil.shape)

# ConvTranspose2d (learned)
t_conv = nn.ConvTranspose2d(8, 4, kernel_size=2, stride=2)
up_ct = t_conv(x)
print('ConvTranspose:', up_ct.shape)

# PixelShuffle r=2: need (B, C*r^2, H, W) = (B, 8*4=32, 4, 4)
x_ps = torch.randn(2, 32, 4, 4)
up_ps = nn.PixelShuffle(2)(x_ps)
print('PixelShuffle: ', up_ps.shape)

21. Where does each piece belong: Lesson 110: Semantic Segmentation & U-Net

Sorting

22. Something is wrong here: checkerboard artifacts from ConvTranspose2d

Anomaly

23. Trap: checkerboard artifacts from ConvTranspose2d

Trap

24. Which of these survive contact with Lesson 110: Semantic Segmentation & U-Net?

Two truths and a lie

25. Segmentation losses — CE, Dice, Focal

Section

Part 4 of 4

26. Why pixel-wise CE fails on imbalanced masks

Concept

Pixel-wise CrossEntropyLoss treats every pixel equally. In medical/satellite images, the foreground class may be 1% of pixels — predicting background everywhere gives 99% accuracy while completely missing the object.

Numbers above are from TinyUNet verification on a (2, 2, 32, 32) mask where only 1/2048 pixels are class-1 — CE barely moves while Dice correctly penalizes the missed minority.

27. Dice loss — overlap-based

Concept

The Dice coefficient (also called F1 score for sets) measures overlap between predicted and ground-truth masks. Dice loss = 1 − Dice, averaged over classes.

\[ \text{Dice}(P, G) = \frac{2 \sum_{i} p_i g_i}{\sum_i p_i + \sum_i g_i + \varepsilon} \qquad \text{Dice loss} = 1 - \frac{1}{C}\sum_{c=1}^C \text{Dice}_c \]

P is the softmax probability map, G is the one-hot ground-truth map, ε prevents division by zero. Range: 0 (perfect) to 1 (no overlap). Naturally handles class imbalance because small classes have the same weight as large ones.

28. By analogy: Dice loss — overlap-based

Analogy

29. Dice loss from scratch

Worked example

Implement Dice loss and verify on random and near-perfect logits. Predict whether the perfect case reaches exactly 0.

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

def dice_loss(pred, target, eps=1e-6):
    """pred: (B, C, H, W) logits; target: (B, H, W) int labels"""
    B, C, H, W = pred.shape
    pred_soft = F.softmax(pred, dim=1)           # (B, C, H, W)
    target_oh = F.one_hot(target, num_classes=C) # (B, H, W, C)
    target_oh = target_oh.permute(0, 3, 1, 2).float()
    intersection = (pred_soft * target_oh).sum(dim=[0, 2, 3])
    denom        = (pred_soft + target_oh).sum(dim=[0, 2, 3])
    dice_per_cls = (2 * intersection + eps) / (denom + eps)
    return 1.0 - dice_per_cls.mean()

torch.manual_seed(7)
B, C, H, W = 2, 3, 8, 8
pred   = torch.randn(B, C, H, W)
target = torch.randint(0, C, (B, H, W))
print('random:', round(dice_loss(pred, target).item(), 4))
perf = F.one_hot(target, C).permute(0,3,1,2).float() * 10
print('perfect:', round(dice_loss(perf, target).item(), 6))

30. Fill in: dice_loss for Dice loss from scratch

Comparison

31. Focal loss — hard-example mining

Concept

Focal loss (Lin 2017, RetinaNet) down-weights easy examples by multiplying the CE by (1 − p_t)^γ where p_t is the predicted probability for the correct class. High confidence → small weight; hard examples → full weight.

\[ \text{FL}(p_t) = -(1-p_t)^\gamma \log(p_t) \qquad \gamma=0 \Rightarrow \text{standard CE} \]

32. Focal loss from scratch

Worked example

Implement focal_loss(gamma=2) and verify it is strictly larger than CE on random logits (easy + hard mixed), and falls near 0 on near-perfect logits. Explain the direction.

import torch, torch.nn.functional as F

def focal_loss(pred, target, gamma=2.0):
    """pred: (B,C,H,W) logits; target: (B,H,W) int labels"""
    B, C, H, W = pred.shape
    log_p = F.log_softmax(pred, dim=1)          # (B,C,H,W)
    p     = log_p.exp()
    log_pt = log_p.gather(1, target.unsqueeze(1)).squeeze(1)  # (B,H,W)
    pt     = p.gather(1, target.unsqueeze(1)).squeeze(1)
    return -(((1 - pt) ** gamma) * log_pt).mean()

torch.manual_seed(7)
B, C, H, W = 2, 3, 8, 8
pred_rnd = torch.randn(B, C, H, W)
target   = torch.randint(0, C, (B, H, W))
ce_rnd   = F.cross_entropy(pred_rnd, target)
fl_rnd   = focal_loss(pred_rnd, target, gamma=2)
print(f'CE random:     {ce_rnd.item():.4f}')
print(f'Focal(g=2) rnd:{fl_rnd.item():.4f}')
perf = F.one_hot(target,C).permute(0,3,1,2).float()*20
print(f'Focal(g=2) perfect: {focal_loss(perf, target).item():.6f}')

33. Without one step: The semantic segmentation recipe

Constraint

34. The semantic segmentation recipe

Pattern

1. Output shape: (B, n_classes, H, W) — one logit map per class, same spatial size as input
2. Encoder: repeated DoubleConv + MaxPool; channels grow, H/W shrink by 2 each stage
3. Bottleneck: deepest DoubleConv — maximum context, minimum H/W
4. Decoder: ConvTranspose2d(k=2,s=2) doubles H/W; then torch.cat([up, enc_skip], dim=1) + DoubleConv
5. Skip connections: concat encoder feature map to matching decoder level (doubles channels); never add
6. Output: Conv2d(base, n_classes, kernel_size=1) — 1×1 conv maps channels to class logits
7. Loss: pixel-wise CE for balanced classes; Dice (overlap) or Focal (hard-example) for imbalanced masks; Dice + CE combined is common
8. ConvTranspose2d size: H_out = (H_in - 1)*stride - 2*padding + kernel_size

35. Where does it stop working: The semantic segmentation recipe

Edge cases

36. Rule out three: Check yourself — skip connection type

Elimination

37. Check yourself — skip connection type

Check

Identify the structural difference before clicking.

38. Answer it before you see the options: Check yourself — ConvTranspose2d output…

Prediction

39. Check yourself — ConvTranspose2d output size

Check

Apply the formula before running.

40. Rule out three: Check yourself — Dice vs CE for imbalanced masks

Elimination

41. Check yourself — Dice vs CE for imbalanced masks

Check

Reason from the loss formulas.

42. Your turn: TinyUNet from scratch

Section

Project

43. Project: TinyUNet on synthetic masks

Concept

Build TinyUNet (1 input channel, 2 classes, base=16) and train it with combined Dice + CE loss on a synthetic segmentation dataset. Three milestones: shape trace → implement Dice loss → train the model.

Build rules: type every line; print shapes after every concat; check logits.shape[-2:] == x.shape[-2:] before computing loss; verify Dice on a trivial case before plugging into the training loop.

44. Break it if you can: Project: TinyUNet on synthetic masks

Counterexample

45. Milestone 1 — TinyUNet shape trace

Worked example

Your turn: implement TinyUNet and verify the output shape. Before coding, predict: after the two ConvTranspose2d upsample steps, does the spatial size return to 32×32?

Hint: use two encoder levels (base=16, base2=32) and one bottleneck (base4=64). Each up step: ConvTranspose2d(in, out, k=2, s=2) then torch.cat([up, enc_skip], dim=1) then DoubleConv.

import torch, torch.nn as nn

class DoubleConv(nn.Module):
    def __init__(self, in_c, out_c):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_c, out_c, 3, padding=1), nn.BatchNorm2d(out_c), nn.ReLU(inplace=True),
            nn.Conv2d(out_c, out_c, 3, padding=1), nn.BatchNorm2d(out_c), nn.ReLU(inplace=True),
        )
    def forward(self, x): return self.conv(x)

class TinyUNet(nn.Module):
    def __init__(self, in_c=1, n_classes=2, base=16):
        super().__init__()
        self.enc1 = DoubleConv(in_c, base)
        self.pool1 = nn.MaxPool2d(2)
        self.enc2 = DoubleConv(base, base*2)
        self.pool2 = nn.MaxPool2d(2)
        self.bottleneck = DoubleConv(base*2, base*4)
        self.up2  = nn.ConvTranspose2d(base*4, base*2, kernel_size=2, stride=2)
        self.dec2 = DoubleConv(base*4, base*2)    # cat: base*2 + base*2
        self.up1  = nn.ConvTranspose2d(base*2, base, kernel_size=2, stride=2)
        self.dec1 = DoubleConv(base*2, base)       # cat: base + base
        self.out_conv = nn.Conv2d(base, n_classes, kernel_size=1)
    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool1(e1))
        b  = self.bottleneck(self.pool2(e2))
        d2 = self.dec2(torch.cat([self.up2(b), e2], dim=1))
        d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
        return self.out_conv(d1)

torch.manual_seed(42)
model = TinyUNet()
x = torch.randn(2, 1, 32, 32)
with torch.no_grad(): out = model(x)
print('params:', sum(p.numel() for p in model.parameters()))
print('output:', out.shape)

46. Milestone 2 — Dice loss verified

Worked example

Your turn: implement dice_loss. Before coding, predict: does Dice = 1 when intersection = 0? What happens when pred_soft is uniform (all classes equally likely)?

Hint: F.softmax(pred, dim=1) gives probabilities; F.one_hot(target, C).permute(0,3,1,2).float() gives the ground-truth mask. Sum intersection and union over [0, 2, 3] (batch + spatial), then take (2*intersection + eps) / (union + eps) per class.

import torch, torch.nn.functional as F

def dice_loss(pred, target, eps=1e-6):
    B, C, H, W = pred.shape
    pred_soft = F.softmax(pred, dim=1)
    target_oh = F.one_hot(target, C).permute(0, 3, 1, 2).float()
    inter = (pred_soft * target_oh).sum(dim=[0, 2, 3])
    denom = (pred_soft + target_oh).sum(dim=[0, 2, 3])
    return 1.0 - ((2 * inter + eps) / (denom + eps)).mean()

torch.manual_seed(7)
pred = torch.randn(2, 3, 8, 8)
tgt  = torch.randint(0, 3, (2, 8, 8))
perf = F.one_hot(tgt, 3).permute(0,3,1,2).float() * 10
print('random:', round(dice_loss(pred, tgt).item(), 4))
print('perfect:', round(dice_loss(perf, tgt).item(), 6))

47. Milestone 3 — train with Dice + CE

Worked example

Your turn: train TinyUNet for 5 epochs with total_loss = CE + Dice. Before running, predict: should total loss start above 1.0 and decrease each epoch?

Hint: create a fixed synthetic target torch.zeros(B, H, W, dtype=torch.long) (all background). CE and Dice together give complementary gradients — CE from per-pixel signal, Dice from overlap.

import torch, torch.nn as nn

torch.manual_seed(42)
model = TinyUNet()   # defined above
opt   = torch.optim.Adam(model.parameters(), lr=1e-3)
ce_fn = nn.CrossEntropyLoss()
x      = torch.randn(2, 1, 32, 32)
target = torch.zeros(2, 32, 32, dtype=torch.long)

for ep in range(5):
    opt.zero_grad()
    out = model(x)
    loss = ce_fn(out, target) + dice_loss(out, target)
    loss.backward(); opt.step()
    print(f'epoch {ep}: total={loss.item():.4f}')

48. What each one costs: Milestone 3 — train with Dice + CE

Trade off

49. The full program

Concept

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

class DoubleConv(nn.Module):
    def __init__(self, i, o):
        super().__init__()
        self.c = nn.Sequential(
            nn.Conv2d(i,o,3,padding=1), nn.BatchNorm2d(o), nn.ReLU(True),
            nn.Conv2d(o,o,3,padding=1), nn.BatchNorm2d(o), nn.ReLU(True))
    def forward(self, x): return self.c(x)

class TinyUNet(nn.Module):
    def __init__(self, in_c=1, n_cls=2, b=16):
        super().__init__()
        self.e1=DoubleConv(in_c,b); self.p1=nn.MaxPool2d(2)
        self.e2=DoubleConv(b,b*2); self.p2=nn.MaxPool2d(2)
        self.bn=DoubleConv(b*2,b*4)
        self.u2=nn.ConvTranspose2d(b*4,b*2,2,2); self.d2=DoubleConv(b*4,b*2)
        self.u1=nn.ConvTranspose2d(b*2,b,2,2);   self.d1=DoubleConv(b*2,b)
        self.out=nn.Conv2d(b,n_cls,1)
    def forward(self,x):
        e1=self.e1(x); e2=self.e2(self.p1(e1)); bn=self.bn(self.p2(e2))
        d2=self.d2(torch.cat([self.u2(bn),e2],1))
        return self.out(self.d1(torch.cat([self.u1(d2),e1],1)))

def dice_loss(pred, target, eps=1e-6):
    B,C,H,W=pred.shape
    ps=F.softmax(pred,1); oh=F.one_hot(target,C).permute(0,3,1,2).float()
    i=(ps*oh).sum([0,2,3]); d=(ps+oh).sum([0,2,3])
    return 1-((2*i+eps)/(d+eps)).mean()

torch.manual_seed(42)
m=TinyUNet(); opt=torch.optim.Adam(m.parameters(),lr=1e-3); ce=nn.CrossEntropyLoss()
x=torch.randn(2,1,32,32); t=torch.zeros(2,32,32,dtype=torch.long)
for ep in range(5):
    opt.zero_grad(); o=m(x)
    loss=ce(o,t)+dice_loss(o,t); loss.backward(); opt.step()
    print(f'ep{ep} total={loss.item():.4f}')
print('output shape:', m(x).shape)  # torch.Size([2, 2, 32, 32])

50. Fill in: value for The full program

Comparison

51. Show it off

Concept

Out loud, slides closed: (1) trace the full TinyUNet forward pass from (2,1,32,32) input to (2,2,32,32) output, naming every shape at each encoder/bottleneck/decoder step; (2) explain why U-Net concatenates skip connections instead of adding them; (3) explain why Dice loss gives 0.50 when a model predicts all-background on a 1/2048 foreground mask, while CE gives only 0.009.

Stretch (homework): replace ConvTranspose2d with bilinear + Conv2d and compare outputs for checkerboard artifacts. Add Focal loss (γ=2) to the combined loss and compare convergence on an imbalanced mask (1% foreground). Try training on sklearn.datasets.load_digits reshaped as segmentation targets.

52. Connect it up: Lesson 110: Semantic Segmentation & U-Net

Connect it up

53. What you can do now

Recap

• Define semantic segmentation output shape (B, n_classes, H, W) and distinguish it from classification
• Trace the U-Net encoder-decoder with skip connections: concat (not add) doubles channels at each decoder level
• Apply the ConvTranspose2d output formula: H_out = (H_in - 1)*stride + kernel_size
• Implement Dice loss from scratch and explain its imbalance robustness (per-class overlap, not per-pixel accuracy)
• Explain when to add Focal loss (γ=2) as a hard-example mining supplement to CE