How to Implement Semantic Segmentation Using the U-Net Architecture: A Complete PyTorch & TensorFlow Guide

Semantic segmentation assigns a class label to every pixel in an image using an encoder-decoder network with skip connections, and the U-Net architecture is the most widely adopted approach for this task.

If you're learning computer vision, understanding how to implement semantic segmentation using the U-Net architecture is a fundamental skill. The microsoft/AI-For-Beginners repository provides production-ready notebooks that walk through the entire pipeline from data loading to model deployment. This guide distills those materials into actionable, runnable code you can adapt for medical imaging, autonomous driving, or any pixel-level prediction task.

What Is Semantic Segmentation and Why U-Net?

Semantic segmentation differs from image classification (one label per image) and object detection (bounding boxes). Instead, it produces a dense prediction map where each pixel receives a class label.

The U-Net architecture, introduced by Ronneberger et al. in U-Net: Convolutional Networks for Biomedical Image Segmentation, solves a critical problem: standard encoder-decoder networks lose fine-grained spatial details during downsampling. U-Net preserves these details through skip connections that concatenate encoder features with decoder features at matching resolutions, creating the characteristic "U" shape.

Repository Structure and Key Files

The AI-For-Beginners course organizes segmentation materials at lessons/4-ComputerVision/12-Segmentation/. Here are the essential files:

File	Purpose
lessons/4-ComputerVision/12-Segmentation/README.md	Theory overview and navigation
SemanticSegmentationPytorch.ipynb	Full PyTorch implementation
SemanticSegmentationTF.ipynb	TensorFlow/Keras equivalent
environment.yml	Dependency specification

Both frameworks follow identical logical steps, so you can choose based on your preference or existing codebase.

Step 1: Data Preprocessing for Segmentation

Before building the model, you need paired image-mask datasets. The notebooks typically use Oxford-IIIT Pets or synthetic binary masks with these preprocessing steps:

• Resize images and masks to a consistent shape (128×128 or 256×256)
• Normalize pixel values to [0, 1] or [-1, 1]
• Apply data augmentation: random horizontal/vertical flips, rotations, and elastic deformations

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os

class SegmentationDataset(Dataset):
    def __init__(self, image_dir, mask_dir, transform=None):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.transform = transform
        self.images = os.listdir(image_dir)
    
    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, idx):
        img_path = os.path.join(self.image_dir, self.images[idx])
        mask_path = os.path.join(self.mask_dir, self.images[idx])
        
        image = Image.open(img_path).convert("RGB")
        mask = Image.open(mask_path).convert("L")  # Grayscale mask

        
        if self.transform:
            image = self.transform(image)
            mask = self.transform(mask)
        
        return image, mask

# Standard preprocessing pipeline

transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
])

Step 2: Building the U-Net Model in PyTorch

The SemanticSegmentationPytorch.ipynb notebook implements U-Net with modular components. Here's the complete architecture following the repository's design patterns:

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

class ConvBlock(nn.Module):
    """
    Two 3×3 convolution layers with ReLU and BatchNorm.
    Replicates the pattern in the official notebook.
    """
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )
    
    def forward(self, x):
        return self.conv(x)

class UNet(nn.Module):
    """
    U-Net implementation for semantic segmentation.
    Encoder path contracts spatial dimensions while increasing channels.
    Decoder path expands spatial dimensions while decreasing channels.
    Skip connections concatenate encoder and decoder features.
    """
    def __init__(self, n_classes=1, in_channels=3, base_channels=64):
        super().__init__()
        
        # Encoder (contracting path)

        self.enc1 = ConvBlock(in_channels, base_channels)
        self.enc2 = ConvBlock(base_channels, base_channels * 2)
        self.enc3 = ConvBlock(base_channels * 2, base_channels * 4)
        self.enc4 = ConvBlock(base_channels * 4, base_channels * 8)
        
        # Bottleneck (deepest layer)

        self.bottleneck = ConvBlock(base_channels * 8, base_channels * 16)
        
        # Decoder (expanding path) with transposed convolutions for upsampling

        self.upconv4 = nn.ConvTranspose2d(
            base_channels * 16, base_channels * 8, kernel_size=2, stride=2
        )
        self.dec4 = ConvBlock(base_channels * 16, base_channels * 8)
        
        self.upconv3 = nn.ConvTranspose2d(
            base_channels * 8, base_channels * 4, kernel_size=2, stride=2
        )
        self.dec3 = ConvBlock(base_channels * 8, base_channels * 4)
        
        self.upconv2 = nn.ConvTranspose2d(
            base_channels * 4, base_channels * 2, kernel_size=2, stride=2
        )
        self.dec2 = ConvBlock(base_channels * 4, base_channels * 2)
        
        self.upconv1 = nn.ConvTranspose2d(
            base_channels * 2, base_channels, kernel_size=2, stride=2
        )
        self.dec1 = ConvBlock(base_channels * 2, base_channels)
        
        # Final 1×1 convolution maps to class labels

        self.out_conv = nn.Conv2d(base_channels, n_classes, kernel_size=1)
    
    def forward(self, x):
        # Encoder with max-pooling between blocks

        e1 = self.enc1(x)
        e2 = self.enc2(F.max_pool2d(e1, kernel_size=2, stride=2))
        e3 = self.enc3(F.max_pool2d(e2, kernel_size=2, stride=2))
        e4 = self.enc4(F.max_pool2d(e3, kernel_size=2, stride=2))
        
        # Bottleneck

        b = self.bottleneck(F.max_pool2d(e4, kernel_size=2, stride=2))
        
        # Decoder with skip connections (concatenate encoder features)

        d4 = self.dec4(torch.cat([self.upconv4(b), e4], dim=1))
        d3 = self.dec3(torch.cat([self.upconv3(d4), e3], dim=1))
        d2 = self.dec2(torch.cat([self.upconv2(d3), e2], dim=1))
        d1 = self.dec1(torch.cat([self.upconv1(d2), e1], dim=1))
        
        # Binary segmentation: sigmoid activation

        # Multi-class: remove sigmoid, use CrossEntropyLoss instead

        return torch.sigmoid(self.out_conv(d1))

Architecture Breakdown

Component	Function	Spatial/Channel Change
Encoder blocks	Feature extraction with increasing abstraction	256×256×3 → 16×16×512
Bottleneck	Deepest feature representation	16×16×512 → 16×16×1024
Decoder blocks	Spatial upsampling with detail recovery	16×16×1024 → 256×256×64
Skip connections	Preserve high-resolution features from encoder	Concatenate matching levels
Output layer	Per-pixel class prediction	256×256×64 → 256×256×C

Step 3: Loss Functions and Metrics for Segmentation

Standard classification losses fail for segmentation due to class imbalance (most pixels are background). The notebooks implement two complementary approaches:

Combined Loss: Dice + Binary Cross-Entropy

class DiceLoss(nn.Module):
    """
    Dice coefficient loss for handling class imbalance.
    Measures overlap between predicted and ground truth masks.
    """
    def __init__(self, smooth=1.0):
        super().__init__()
        self.smooth = smooth
    
    def forward(self, predictions, targets):
        predictions = predictions.view(-1)
        targets = targets.view(-1)
        
        intersection = (predictions * targets).sum()
        dice = (2. * intersection + self.smooth) / (
            predictions.sum() + targets.sum() + self.smooth
        )
        return 1 - dice

class CombinedLoss(nn.Module):
    """
    Weighted combination of Dice and BCE losses.
    Used in the official PyTorch notebook.
    """
    def __init__(self, dice_weight=0.5, bce_weight=0.5):
        super().__init__()
        self.dice = DiceLoss()
        self.bce = nn.BCELoss()
        self.dice_weight = dice_weight
        self.bce_weight = bce_weight
    
    def forward(self, pred, target):
        return self.dice_weight * self.dice(pred, target) + \
               self.bce_weight * self.bce(pred, target)

Intersection over Union (IoU) Metric

def iou_score(prediction, target, threshold=0.5):
    """
    Calculate Intersection over Union for binary segmentation.
    """
    prediction = (prediction > threshold).float()
    target = target.float()
    
    intersection = (prediction * target).sum()
    union = prediction.sum() + target.sum() - intersection
    
    return (intersection + 1e-6) / (union + 1e-6)

Step 4: Training Loop Implementation

The SemanticSegmentationPytorch.ipynb notebook uses standard PyTorch training patterns with device management:

def train_segmentation_model(
    model, 
    train_loader, 
    val_loader, 
    num_epochs=25,
    learning_rate=0.001,
    device='cuda'
):
    model = model.to(device)
    criterion = CombinedLoss(dice_weight=0.5, bce_weight=0.5)
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='max', factor=0.5, patience=5
    )
    
    best_iou = 0.0
    
    for epoch in range(num_epochs):
        # Training phase

        model.train()
        epoch_loss = 0.0
        
        for images, masks in train_loader:
            images = images.to(device)
            masks = masks.to(device)
            
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, masks)
            loss.backward()
            optimizer.step()
            
            epoch_loss += loss.item()
        
        # Validation phase

        model.eval()
        val_iou = 0.0
        
        with torch.no_grad():
            for images, masks in val_loader:
                images = images.to(device)
                masks = masks.to(device)
                
                outputs = model(images)
                val_iou += iou_score(outputs, masks).item()
        
        avg_val_iou = val_iou / len(val_loader)
        scheduler.step(avg_val_iou)
        
        print(f"Epoch {epoch+1}/{num_epochs}: "
              f"Loss={epoch_loss/len(train_loader):.4f}, "
              f"Val IoU={avg_val_iou:.4f}")
        
        # Save best model

        if avg_val_iou > best_iou:
            best_iou = avg_val_iou
            torch.save(model.state_dict(), 'best_unet_model.pth')
    
    return model

Step 5: Inference and Visualization

After training, generate predictions and overlay masks on original images:

import matplotlib.pyplot as plt

def visualize_predictions(model, dataloader, device='cuda', num_samples=3):
    model.eval()
    model = model.to(device)
    
    images, masks = next(iter(dataloader))
    images = images[:num_samples].to(device)
    masks = masks[:num_samples]
    
    with torch.no_grad():
        predictions = model(images)
        predictions = (predictions > 0.5).float().cpu()
    
    fig, axes = plt.subplots(num_samples, 3, figsize=(12, 4*num_samples))
    
    for i in range(num_samples):
        # Original image

        axes[i, 0].imshow(images[i].cpu().permute(1, 2, 0))
        axes[i, 0].set_title('Input Image')
        axes[i, 0].axis('off')
        
        # Ground truth mask

        axes[i, 1].imshow(masks[i].squeeze(), cmap='gray')
        axes[i, 1].set_title('Ground Truth')
        axes[i, 1].axis('off')
        
        # Predicted mask

        axes[i, 2].imshow(predictions[i].squeeze(), cmap='gray')
        axes[i, 2].set_title(f'Prediction (IoU: {iou_score(predictions[i], masks[i]):.3f})')
        axes[i, 2].axis('off')
    
    plt.tight_layout()
    return fig

TensorFlow/Keras Implementation

The SemanticSegmentationTF.ipynb notebook provides equivalent functionality. Key differences:

• Use tf.keras.layers.Conv2DTranspose instead of ConvTranspose2d
• Implement skip connections with tf.keras.layers.Concatenate
• Use tf.image for preprocessing instead of torchvision

The overall architecture—encoder, bottleneck, decoder with skip connections—remains identical, making it straightforward to port models between frameworks.

Advanced: Swapping the Encoder Backbone

The notebooks demonstrate how to replace the standard encoder with pretrained backbones like ResNet-50 or EfficientNet. This transfer learning approach:

1. Load pretrained weights from torchvision.models or tf.keras.applications
2. Remove the final classification layers
3. Connect the feature pyramids to the U-Net decoder
4. Fine-tune with a lower learning rate for the backbone

This typically improves IoU by 5-15 points on limited datasets.

Summary

• Semantic segmentation requires per-pixel classification, and U-Net's skip connections preserve the spatial precision needed for accurate masks
• The microsoft/AI-For-Beginners repository provides complete, beginner-friendly implementations in both PyTorch (SemanticSegmentationPytorch.ipynb) and TensorFlow (SemanticSegmentationTF.ipynb)
• A proper implementation includes: modular ConvBlock layers, symmetric encoder-decoder paths, skip connections via concatenation, combined Dice/BCE loss, and IoU monitoring
• For production use, consider pretrained encoder backbones and data augmentation to improve generalization

Frequently Asked Questions

What is the difference between U-Net and standard encoder-decoder networks?

U-Net adds skip connections that concatenate encoder features with decoder features at the same spatial resolution. Standard encoder-decoder networks only pass information through the bottleneck, losing fine details. These skip connections allow U-Net to segment small, precise structures like cell boundaries or road edges.

How do I handle multi-class segmentation with U-Net?

Change n_classes to your number of classes and remove the final sigmoid activation. Use nn.CrossEntropyLoss (PyTorch) or SparseCategoricalCrossentropy (TensorFlow) as your loss function. The output shape becomes [batch, n_classes, height, width], and you apply argmax across the channel dimension for final predictions.

Why does my U-Net produce blurry segmentation masks?

Blurry outputs typically indicate: (1) insufficient model capacity—try increasing base_channels; (2) missing skip connections—verify concatenation operations; (3) over-regularization—reduce dropout or weight decay; or (4) training time—segmentation requires more epochs than classification. Check that your upsampling uses transposed convolutions rather than simple interpolation.

Can I use U-Net for 3D volumetric segmentation?

Yes. Replace all 2D convolutions with 3D convolutions (nn.Conv3d), 2D pooling with 3D pooling, and ensure your data loader provides volumetric patches. The architecture logic remains identical, though memory requirements increase substantially—consider using a smaller base_channels or patch-based training for large volumes.

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