Semantic Segmentation with U-Net in Computer Vision Lessons: A Complete PyTorch Implementation

The Microsoft AI-For-Beginners curriculum teaches semantic segmentation using a custom U-Net implementation with skip connections that preserves spatial detail by concatenating encoder feature maps with upsampled decoder outputs, achieving pixel-level classification on the PH² dermoscopy dataset.

Semantic segmentation with U-Net transforms image classification into a dense prediction task where every pixel receives a class label. The microsoft/AI-For-Beginners repository provides a complete educational implementation in Lesson 4-ComputerVision/12-Segmentation, demonstrating how to build this encoder-decoder architecture from scratch using PyTorch.

Understanding Semantic Segmentation in the Curriculum

What is Semantic Segmentation?

Unlike object detection that draws bounding boxes, semantic segmentation assigns a class label to every pixel in an image. This dense classification problem requires architectures that preserve spatial resolution while extracting hierarchical features.

Lesson Structure and Resources

The lesson is located in lessons/4-ComputerVision/12-Segmentation/ and contains:

• README.md – Conceptual overview and architecture diagrams
• SemanticSegmentationPytorch.ipynb – Complete PyTorch implementation including the UNet class and training loop
• SemanticSegmentationTF.ipynb – TensorFlow/Keras alternative implementation

The curriculum compares a simple encoder-decoder SegNet baseline against the more accurate U-Net architecture that utilizes skip connections.

The U-Net Architecture Design

Skip Connections and Spatial Preservation

The primary limitation of vanilla encoder-decoder models is the loss of spatial detail caused by repeated downsampling. Inspired by the original 2015 paper U-Net: Convolutional Networks for Biomedical Image Segmentation by Ronneberger et al., the architecture introduces skip connections that directly concatenate feature maps from each encoder stage with their corresponding upsampled decoder counterparts.

This design preserves high-resolution spatial information while still benefiting from deep hierarchical features extracted by the encoder.

Encoder-Decoder Channel Dimensions

The implementation in SemanticSegmentationPytorch.ipynb uses specific channel progressions:

Encoder pathway (contracting):

• Input: 3 channels (RGB)
• Block 0: 16 channels
• Block 1: 32 channels
• Block 2: 64 channels
• Block 3: 128 channels
• Bottleneck: 256 channels

Decoder pathway (expanding):

• Level 0: 384 channels (256 upsampled + 128 skip) → 128 output
• Level 1: 192 channels (128 upsampled + 64 skip) → 64 output
• Level 2: 96 channels (64 upsampled + 32 skip) → 32 output
• Level 3: 48 channels (32 upsampled + 16 skip) → 1 output (mask)

PyTorch Implementation Details

UNet Class Definition

The UNet class defined in SemanticSegmentationPytorch.ipynb (around lines 73-90) implements the complete architecture using standard PyTorch modules:

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

class UNet(nn.Module):
    def __init__(self):
        super().__init__()
        # Encoder blocks: Conv2d → ReLU → BatchNorm → MaxPool2d

        self.enc_conv0 = nn.Conv2d(3, 16, 3, padding=1)
        self.bn0 = nn.BatchNorm2d(16)
        self.pool0 = nn.MaxPool2d(2)
        
        self.enc_conv1 = nn.Conv2d(16, 32, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.pool1 = nn.MaxPool2d(2)
        
        self.enc_conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.pool2 = nn.MaxPool2d(2)
        
        self.enc_conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.pool3 = nn.MaxPool2d(2)
        
        # Bottleneck

        self.bottleneck = nn.Conv2d(128, 256, 3, padding=1)
        
        # Decoder blocks with skip connections

        self.upsample0 = nn.UpsamplingBilinear2d(2)
        self.dec_conv0 = nn.Conv2d(384, 128, 3, padding=1)  # 256 + 128

        self.dec_bn0 = nn.BatchNorm2d(128)
        
        self.upsample1 = nn.UpsamplingBilinear2d(2)
        self.dec_conv1 = nn.Conv2d(192, 64, 3, padding=1)   # 128 + 64

        self.dec_bn1 = nn.BatchNorm2d(64)
        
        self.upsample2 = nn.UpsamplingBilinear2d(2)
        self.dec_conv2 = nn.Conv2d(96, 32, 3, padding=1)   # 64 + 32

        self.dec_bn2 = nn.BatchNorm2d(32)
        
        self.upsample3 = nn.UpsamplingBilinear2d(2)
        self.dec_conv3 = nn.Conv2d(48, 1, 1)               # 32 + 16 → 1-channel

        self.sigmoid = nn.Sigmoid()

Forward Pass with Concatenation

The forward method implements the skip connections using torch.cat to concatenate encoder features with upsampled decoder features along the channel dimension:

    def forward(self, x):
        # Encoder forward pass

        e0 = self.pool0(self.bn0(F.relu(self.enc_conv0(x))))
        e1 = self.pool1(self.bn1(F.relu(self.enc_conv1(e0))))
        e2 = self.pool2(self.bn2(F.relu(self.enc_conv2(e1))))
        e3 = self.pool3(self.bn3(F.relu(self.enc_conv3(e2))))
        
        # Bottleneck

        b = self.bottleneck(e3)
        
        # Decoder with skip connections (concatenate along dim=1)

        d0 = self.dec_bn0(F.relu(self.dec_conv0(
            torch.cat((self.upsample0(b), e3), dim=1))))
        d1 = self.dec_bn1(F.relu(self.dec_conv1(
            torch.cat((self.upsample1(d0), e2), dim=1))))
        d2 = self.dec_bn2(F.relu(self.dec_conv2(
            torch.cat((self.upsample2(d1), e1), dim=1))))
        d3 = self.sigmoid(self.dec_conv3(
            torch.cat((self.upsample3(d2), e0), dim=1)))
        return d3

Training Configuration and Dataset

PH² Dermoscopy Dataset

The notebook uses the PH² dermoscopy dataset for binary segmentation tasks. The data loading utilities (cells 5-7 in the notebook) split the dataset into 90% training and 10% testing:

from torch.utils.data import DataLoader

train_dataset, test_dataset = load_dataset(train_part=0.9)
train_loader = DataLoader(list(zip(train_dataset[0], train_dataset[1])),
                          batch_size=32, shuffle=True)
test_loader = DataLoader(list(zip(test_dataset[0], test_dataset[1])),
                         batch_size=1, shuffle=False)

Loss Function and Optimization

For the binary mask prediction task, the implementation uses nn.BCEWithLogitsLoss (defined around lines 60-66), which combines a sigmoid layer and binary cross-entropy loss for numerical stability:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = UNet().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-6)
criterion = nn.BCEWithLogitsLoss()

# Training loop excerpt (lines 90-120 in notebook)

for epoch in range(30):
    model.train()
    for imgs, masks in train_loader:
        imgs, masks = imgs.to(device), masks.to(device)
        preds = model(imgs)
        loss = criterion(preds, masks)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

The training loop reports train and test loss per epoch, followed by visualization cells that compare predicted masks against ground truth.

Summary

• Semantic segmentation assigns class labels to every pixel, implemented in lessons/4-ComputerVision/12-Segmentation/ of the AI-For-Beginners repository.
• The U-Net architecture uses skip connections to preserve spatial detail by concatenating encoder feature maps with decoder upsampling outputs.
• Channel dimensions progress from 16→32→64→128→256 in the encoder and 384→192→96→48 in the decoder due to concatenation.
• The implementation uses PyTorch modules including Conv2d, BatchNorm2d, MaxPool2d, and UpsamplingBilinear2d.
• Binary cross-entropy with logits (nn.BCEWithLogitsLoss) provides the loss function for training on the PH² dataset.

Frequently Asked Questions

What is the difference between U-Net and SegNet in the lesson?

SegNet uses a standard encoder-decoder architecture without skip connections, relying solely on the bottleneck features for reconstruction. U-Net preserves high-resolution details by concatenating encoder feature maps directly with the corresponding decoder stages, resulting in more precise segmentation boundaries especially for small objects.

Why does the first decoder convolution use 384 input channels?

The first decoder block receives 256 channels from the upsampled bottleneck and 128 channels from the encoder skip connection, which are concatenated along the channel dimension (256 + 128 = 384). This pattern continues throughout the decoder: 128+64=192, 64+32=96, and 32+16=48, enabling the network to combine deep semantic features with shallow spatial details.

What dataset is used for training the segmentation model?

The notebook uses the PH² dermoscopy dataset, a medical imaging dataset for skin lesion analysis. The implementation performs binary segmentation to separate lesion regions from healthy skin, with data utilities provided within SemanticSegmentationPytorch.ipynb to handle the 90/10 train-test split.

How do skip connections improve segmentation accuracy?

Skip connections mitigate the information bottleneck caused by repeated downsampling in the encoder. By concatenating high-resolution encoder features with upsampled decoder features, the network retains fine-grained spatial information (like object boundaries) that would otherwise be lost, while still utilizing the deep semantic understanding captured in the bottleneck.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}