Day 66: Building a tiny U-Net from scratch (part 2): decoder & training
The expanding path and putting it together
The decoder mirrors the encoder in reverse: upsample (transposed convolution or interpolation) to grow spatial size, concatenate the matching encoder skip, then a conv block. A final 1×1 convolution maps to the number of output classes (for a binary person mask, one channel). Trained with a segmentation loss (Dice or binary cross-entropy), you now have a working U-Net you built end to end.
import torch, torch.nn as nn
class Decoder(nn.Module):
def __init__(self):
super().__init__()
self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
self.d2 = conv_block(256, 128) # 128 (upsampled) + 128 (skip s2)
self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
self.d1 = conv_block(128, 64)
self.head = nn.Conv2d(64, 1, 1) # 1 channel: person vs background
def forward(self, x, skips):
s1, s2 = skips
x = self.d2(torch.cat([self.up2(x), s2], dim=1)) # concat skip
x = self.d1(torch.cat([self.up1(x), s1], dim=1))
return self.head(x) # raw logits, HxWDice loss for imbalanced masks
A person occupies a fraction of the image, so background pixels dominate — plain cross-entropy can be lazy. Dice loss directly optimizes mask overlap (it's differentiable IoU-flavored), handling the imbalance better. Combining Dice + BCE is a common, robust choice for segmentation. Same imbalance instinct from Day 20, new domain.
Key terms
- Transposed convolution
- A learnable upsampling operation that increases spatial resolution in the decoder.
- Skip concatenation
- Joining an encoder skip feature map with the upsampled decoder features along the channel dimension.
- Dice loss
- A segmentation loss optimizing overlap between predicted and true masks, robust to class imbalance.
Why is Dice loss often preferred over plain cross-entropy for training a person-segmentation U-Net?
