Semantic Segmentation and U-Net Explained

Semantic segmentation assigns a class to every pixel: a road scene might label road, sky, cars, and pedestrians. In the PyTorch convention used here, the model produces logits shaped \((N,C,H,W)\), where \(N\) is batch size, \(C\) the number of classes, and \(H,W\) the image height and width. Selecting the largest class score at each pixel gives an \((N,H,W)\) label map. The encoder extracts context while the decoder builds a dense prediction, using U-Net’s skip connections to access finer-resolution features.

Three related tasks

Semantic segmentation gives both adjacent cars the same car class, without assigning separate object identities. Instance segmentation assigns an individual mask to each countable object. Panoptic segmentation combines those object identities with semantic labels for regions such as sky and road. Decide whether individual objects need separate identities before defining the annotation format.

Convolution, padding, stride, and pooling are introduced in CNN Fundamentals: From Convolution to Image Classification.

The encoder–decoder shape

An encoder reduces spatial size while extracting features over larger receptive fields. A backbone with output stride 32 maps 224-by-224 inputs to 7-by-7 features, but that reduction is an architectural choice. The compact U-Net below pools three times, so 128 becomes 64, then 32, then 16. Its channel counts grow as spatial size decreases.

The decoder upsamples features and learns how to convert them into per-pixel predictions. Interpolation alone cannot recover image-specific detail absent from its input. Coarse features can still locate broad regions, while higher-resolution encoder features help the decoder refine boundaries. Upsampling increases the number of output positions; it does not undo information loss automatically.

Transposed convolution

For fixed weights and no bias, write a convolution as a matrix operation \(y=Ax\). A transposed convolution applies \(A^\top\) with compatible input and output shapes. This is also the linear operation used to propagate gradients from \(y\) to \(x\); it is generally not \(A^{-1}\). A small example is \(A=[1,2]\): it maps \([3,4]^\top\) to 11, while \(A^\top 11=[11,22]^\top\), which does not recover the input. Transposed convolution is one learnable upsampling option.

Along either spatial axis, \[n_{\mathrm{out}}=(n_{\mathrm{in}}-1)s-2p+d(k-1)+o+1.\] Here \(s\) is stride, \(p\) padding, \(d\) dilation, \(k\) kernel size, and \(o\) output padding. With \(k=2,s=2,p=0,d=1,o=0\), the size doubles. Stride 2 alone does not ensure doubling: \(k=3\) with those other settings gives \(2n+1\). Output padding selects the output size when striding leaves an ambiguity; it does not simply append zero-valued pixels.

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

up = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
print(up(torch.randn(1, 64, 28, 28)).shape)
# torch.Size([1, 32, 56, 56])

resize_up = nn.Sequential(
    nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
    nn.Conv2d(64, 32, kernel_size=3, padding=1),
)
print(resize_up(torch.randn(1, 64, 28, 28)).shape)
# torch.Size([1, 32, 56, 56])

For dilation 1, a kernel size not divisible by the stride can create uneven overlap in the interior: some output positions receive more contributions than others. This can encourage checkerboard patterns. A divisible kernel size removes that particular interior overlap imbalance, but learned weights and boundary effects can still create artifacts. Resize-then-convolve avoids this transposed-convolution overlap mechanism; it does not guarantee artifact-free predictions.

Skip connections: the U-Net idea

U-Net combines a contracting path with an expanding path and concatenates encoder features into decoder stages at matching spatial scales. These features give the decoder access to details that need not pass through the bottleneck. They are learned representations, so a skip connection does not guarantee exact boundaries. The code below is a compact U-Net variant with padded convolutions, BatchNorm, and three pooling stages. The original U-Net used unpadded convolutions and cropped encoder features for concatenation.

At each decoder stage, upsampled features and an encoder skip share height and width. Concatenation preserves their separate channels, and subsequent convolutions learn how to combine them. A residual addition instead sums corresponding channels of compatible tensors. Both offer additional paths through a network, but their operations and architectural purposes differ.

def double_conv(c_in, c_out):
    return nn.Sequential(
        nn.Conv2d(c_in, c_out, 3, padding=1, bias=False),
        nn.BatchNorm2d(c_out), nn.ReLU(inplace=True),
        nn.Conv2d(c_out, c_out, 3, padding=1, bias=False),
        nn.BatchNorm2d(c_out), nn.ReLU(inplace=True))

class UNet(nn.Module):
    def __init__(self, in_ch=3, num_classes=2, base=64):
        super().__init__()
        self.enc1 = double_conv(in_ch, base)
        self.enc2 = double_conv(base, base * 2)
        self.enc3 = double_conv(base * 2, base * 4)
        self.bottleneck = double_conv(base * 4, base * 8)
        self.pool = nn.MaxPool2d(2)

        self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, 2)
        self.dec3 = double_conv(base * 8, base * 4)     # 4 + 4 after concat
        self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, 2)
        self.dec2 = double_conv(base * 4, base * 2)
        self.up1 = nn.ConvTranspose2d(base * 2, base, 2, 2)
        self.dec1 = double_conv(base * 2, base)
        self.head = nn.Conv2d(base, num_classes, 1)     # 1x1 to class logits

    def forward(self, x):
        if (x.ndim != 4 or min(x.shape[-2:]) < 8
                or x.shape[-2] % 8 or x.shape[-1] % 8):
            raise ValueError("expected (N,C,H,W) with H,W positive multiples of 8")
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool(e1))
        e3 = self.enc3(self.pool(e2))
        b = self.bottleneck(self.pool(e3))

        d3 = self.dec3(torch.cat([self.up3(b), e3], dim=1))
        d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1))
        d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
        return self.head(d1)                            # (N, C, H, W) logits

model = UNet(3, 2, base=16)
print(model(torch.randn(1, 3, 128, 128)).shape)
# torch.Size([1, 2, 128, 128])

For base=16, the first decoder upsampling produces 64 channels and its encoder skip also has 64; concatenating gives 128 channels for dec3. Matching spatial sizes is required for concatenation, while matching the next convolution’s input-channel count is a separate check. The output contains two logits at each pixel. Use argmax(dim=1) for a hard two-class label map; use softmax across channels if probabilities are needed.

This implementation requires height and width divisible by \(2^3=8\), not 32. Its early check reports incompatible input sizes before concatenation. Padding to the next multiple and cropping the logits back is one option; explicit resizing to skip sizes is another design. Padding changes boundary context, so use a consistent policy for training and inference. BatchNorm in training mode additionally needs more than one value per channel across batch and spatial positions at every stage.

def forward_padded(model, images):
    height, width = images.shape[-2:]
    if min(height, width) <= 0:
        raise ValueError("spatial dimensions must be positive")
    padded = F.pad(images, (0, (-width) % 8, 0, (-height) % 8))
    return model(padded)[..., :height, :width]

model.eval()
with torch.no_grad():
    print(forward_padded(model, torch.randn(1, 3, 127, 131)).shape)
# torch.Size([1, 2, 127, 131])

The helper pads only the right and bottom edges, then removes those positions from the logits. Compute the loss against the original-size mask after this crop. If you instead keep padded logits for training, exclude artificial mask pixels from the loss. The example uses evaluation mode; use model.train() when fitting parameters.

Loss functions for imbalanced masks

Per-pixel cross-entropy is a baseline. Class proportions vary by dataset, but when foreground occupies 0.5% of the pixels, predicting background everywhere gives 99.5% accuracy while missing every foreground pixel. A region-overlap metric helps reveal that failure; the high accuracy does not mean the foreground task is solved.

For hard binary masks, foreground Dice is \(2TP/(2TP+FP+FN)\). Here TP counts correctly predicted foreground pixels, FP counts false foreground predictions, and FN counts missed foreground pixels. Soft Dice replaces predicted membership with foreground probabilities \(p_i\), while \(g_i\) is the binary target. For one image, the following smoothed loss is a differentiable surrogate for overlap, not identical to the Dice of thresholded predictions:

\[\mathcal{L}_{\text{Dice}}=1-\frac{2\sum_i p_i g_i+\epsilon}{\sum_i p_i+\sum_i g_i+\epsilon}\]

def soft_dice_loss(logits, targets, eps=1.0, valid=None):
    """Binary (N,1,H,W) logits and 0/1 targets; optional Boolean validity mask."""
    if logits.ndim != 4 or logits.shape[1] != 1 or targets.shape != logits.shape:
        raise ValueError("logits and targets must have identical (N,1,H,W) shapes")
    if eps <= 0:
        raise ValueError("eps must be positive")
    if valid is None:
        valid = torch.ones_like(targets, dtype=torch.bool)
    if valid.shape != targets.shape or valid.dtype != torch.bool:
        raise ValueError("valid must be a Boolean mask with the same shape")
    probs = torch.where(valid, torch.sigmoid(logits), 0.0)
    target = torch.where(valid, targets.to(logits.dtype), 0.0)
    inter = (probs * target).sum((1, 2, 3))
    denom = probs.sum((1, 2, 3)) + target.sum((1, 2, 3))
    loss = 1 - (2 * inter + eps) / (denom + eps)
    usable = valid.flatten(1).any(dim=1)
    if not usable.any():
        return logits.sum() * 0.0
    return loss[usable].mean()

logits = torch.full((1, 1, 4, 4), 10.0)
targets = torch.ones_like(logits)
print(round(soft_dice_loss(logits, targets).item(), 6))
# 0.000022

The function computes Dice per image and averages over images with at least one valid pixel; pooling the whole batch first defines a different loss. The smoothing constant affects small and empty masks. For an empty target, its numerator is just eps, so predictions still affect the loss. A fully ignored image is excluded. In this example, sigmoid(10) is close to one, so the printed loss is small but not zero.

The two-channel U-Net above fits mutually exclusive background/foreground labels with CrossEntropyLoss, using integer targets shaped (N,H,W). The binary Dice helper instead needs one foreground logit and floating targets shaped (N,1,H,W). Instantiate a one-output model for the next example and combine BCEWithLogitsLoss with Dice. For multiclass Dice, use softmax probabilities and one-hot targets, with an explicit class-averaging and background policy.

torch.manual_seed(0)
binary_model = UNet(in_ch=3, num_classes=1, base=8)
optimizer = torch.optim.SGD(binary_model.parameters(), lr=0.01)
images = torch.randn(2, 3, 32, 32)
targets = torch.zeros(2, 1, 32, 32)
targets[..., 8:24, 8:24] = 1.0

binary_model.train()
optimizer.zero_grad(set_to_none=True)
logits = binary_model(images)
loss = F.binary_cross_entropy_with_logits(logits, targets) + soft_dice_loss(logits, targets)
before = binary_model.head.weight.detach().clone()
loss.backward()
optimizer.step()
print("finite loss:", bool(torch.isfinite(loss)))
print("head updated:", not torch.equal(before, binary_model.head.weight))
# finite loss: True
# head updated: True

binary_model.eval()
with torch.no_grad():
    mask = binary_model(images).sigmoid() >= 0.5
print(tuple(mask.shape))
# (2, 1, 32, 32)

This synthetic batch checks a training step and mask shape; it measures no segmentation quality. Adding BCE and Dice is a candidate objective, with weights to choose on validation data. Their gradient scales depend on reduction, mask size, predictions, and smoothing, so equal coefficients do not imply equal influence. The example uses fully labeled masks. For void pixels, apply the same validity mask to both losses; BCE has no ignore_index. Choose a binary probability threshold on validation data and keep it fixed for testing.

Evaluation

For class \(c\), mask IoU is \(TP_c/(TP_c+FP_c+FN_c)\), using pixel counts. One common dataset-level metric accumulates a confusion matrix across all valid pixels, computes each class’s IoU, and averages the included classes. For a class with nonzero denominator, hard Dice and IoU satisfy \(\mathrm{Dice}=2\mathrm{IoU}/(1+\mathrm{IoU})\). These hard-mask metrics differ from the soft training loss.

State which classes are included, whether background is included, and how zero-union classes are handled. A class absent from both ground truth and predictions has undefined IoU; the example excludes it. A class absent from ground truth but predicted somewhere has positive union and IoU zero. Averaging per-image IoU is another valid statistic when defined consistently, but weights images differently from pooling counts. Report per-class results because the mean can obscure a poorly performing rare class.

For two classes, let rows of a confusion matrix denote truth and columns denote predictions. The matrix below has 90 correctly labeled background pixels, 2 false foreground predictions, 3 missed foreground pixels, and 5 correctly labeled foreground pixels. The foreground IoU is 5/(5+2+3)=0.5, although pixel accuracy is 0.95. These counts could be summed over several images before computing the ratios.

confusion = torch.tensor([[90, 2], [3, 5]], dtype=torch.float64)
tp = confusion.diag()
union = confusion.sum(1) + confusion.sum(0) - tp
present = union > 0
class_iou = torch.full_like(union, float("nan"))
class_iou[present] = tp[present] / union[present]
mean_iou = class_iou[present].mean()
print([round(x, 4) for x in class_iou.tolist()])
print(round(mean_iou.item(), 4))
# [0.9474, 0.5]
# 0.7237

Exclude void pixels before accumulating this matrix. If there are no included classes with positive union, report the metric as undefined rather than forcing a zero. Keep related frames, scans, or tiles together when the intended test case is a new source image or subject; overlapping tiles across splits can make evaluation optimistic.

Boxes, IoU, and non-maximum suppression are covered in Object Detection: Localization, YOLO, IoU, NMS, and R-CNN.

Data pipeline checks

• Apply the same crop coordinates, flips, and other geometric parameters to each image-mask pair. Choose interpolation separately for image values and hard mask IDs; drawing independent random geometry can misalign the pair.
• Resize masks with nearest-neighbor interpolation. Bilinear interpolation mixes numeric IDs as if they were continuous values; casting the resulting fractions to integers can introduce an unrelated class. This rule concerns hard class-ID masks, not probability maps.
• Color jitter and blur apply to the image only.
• For the class-index form of multiclass cross-entropy, masks are integer indices, shape \((N,H,W)\) with no channel dimension, dtype int64.
• Use ignore_index for void regions in class-index cross-entropy. Mask these pixels separately in Dice and metrics as well. Handle an all-void batch explicitly instead of averaging over zero valid pixels.
• Overlay a few masks on their images and look at them before training. This can reveal incorrect class mappings, spatial misalignment, or image/mask flips before a long training run.

For larger images, choose between resizing and tiled inference based on object size and memory limits. Tiling can remove context and create edge seams; if overlapping tiles are blended, combine probabilities or logits consistently before producing the hard mask. Validate the same resize or tiling policy used at deployment.

Inspect predicted masks as well as aggregate scores. A region score may look acceptable while thin structures or boundaries are poor; add a task-appropriate boundary measure when those errors matter.

Exercises

1. Divisibility. A U-Net has three pooling stages. Compute the output shape for inputs of 128, 224, and 250 pixels, and state which size is incompatible and where an unchecked decoder would fail.

You should get: two compatible sizes and one that the implementation rejects before running the encoder.

for n in (128, 224, 250):
    sizes = [n]
    for _ in range(3): sizes.append(sizes[-1] // 2)
    up = [sizes[-1]]
    for _ in range(3): up.append(up[-1] * 2)
    print(n, "down", sizes, "up", up, "| match:", up[1:] == sizes[-2::-1])
# 128 down [128, 64, 32, 16] up [16, 32, 64, 128] | match: True
# 224 down [224, 112, 56, 28] up [28, 56, 112, 224] | match: True
# 250 down [250, 125, 62, 31] up [31, 62, 124, 248] | match: False

2. Foreground overlap and pixel accuracy. A scan has 0.5% foreground pixels. Compute pixel accuracy and Dice for a model that predicts all background. Then compute both for a model that finds half the foreground with no false positives.

You should get: 99.5% accuracy with zero foreground Dice, then 99.75% accuracy with Dice about 0.667.

import numpy as np
gt = np.zeros(10_000); gt[:50] = 1                 # 0.5% foreground
def scores(pred):
    acc = (pred == gt).mean()
    inter = (pred*gt).sum()
    dice = 2*inter / (pred.sum() + gt.sum() + 1e-9)
    return round(acc, 4), round(dice, 4)
print("all background", scores(np.zeros(10_000)))
# all background (0.995, 0.0)
half = np.zeros(10_000); half[:25] = 1
print("half found", scores(half))
# half found (0.9975, 0.6667)

3. Mask interpolation. Resize a segmentation mask containing only class indices 1 and 3, first with bilinear interpolation and then with nearest neighbour. Report the set of unique values in each result.

You should get: invented class indices in one result and only the original two in the other.

import numpy as np, torch, torch.nn.functional as F
mask = torch.zeros(1, 1, 4, 4)
mask[..., :2] = 1.0; mask[..., 2:] = 3.0
bil = F.interpolate(mask, size=(8, 8), mode="bilinear", align_corners=False)
near = F.interpolate(mask, size=(8, 8), mode="nearest")
print("bilinear", sorted(set(bil.flatten().tolist())))
# bilinear [1.0, 1.5, 2.5, 3.0]
print("nearest", sorted(set(near.flatten().tolist())))
# nearest [1.0, 3.0]

References

• Long, Shelhamer, and Darrell (2015). Fully Convolutional Networks for Semantic Segmentation.
• Ronneberger, Fischer, and Brox (2015). U-Net: Convolutional Networks for Biomedical Image Segmentation. MICCAI.
• Milletari, Navab, and Ahmadi (2016). V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation. 3DV.
• Odena, Dumoulin, and Olah (2016). Deconvolution and Checkerboard Artifacts. Distill.
• Chen, Zhu, Papandreou, Schroff, and Adam (2018). Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation. ECCV.
• Kirillov, He, Girshick, Rother, and Dollar (2019). Panoptic Segmentation.
• PyTorch. ConvTranspose2d and CrossEntropyLoss.