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

> Learn semantic segmentation with U-Net in computer vision. This PyTorch implementation from Microsoft AI-For-Beginners preserves spatial detail for pixel-level classification. Get the complete guide.

- Repository: [Microsoft/AI-For-Beginners](https://github.com/microsoft/AI-For-Beginners)
- Tags: tutorial
- Published: 2026-08-05

---

**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`](https://github.com/microsoft/AI-For-Beginners/blob/main/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:

```python
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:

```python
    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:

```python
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:

```python
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.