How to Implement Semantic Segmentation with U-Net Models in PyTorch

Implement semantic segmentation in PyTorch by constructing a U-Net architecture with a ResNet34 encoder, transposed convolutional decoder, and skip connections, then train using binary cross-entropy loss with logits.

Semantic segmentation classifies every pixel in an image into a specific category, enabling applications like medical imaging and autonomous driving. In the microsoft/AI-For-Beginners repository, the notebook SemanticSegmentationPytorch.ipynb demonstrates a complete implementation using a U-Net architecture that combines a pretrained encoder with a upsampling decoder. This guide walks through the actual source code to build, train, and evaluate a pixel-level classification model.

Understanding the U-Net Architecture for Semantic Segmentation

U-Net follows an encoder-decoder structure with skip connections that preserve spatial information lost during downsampling. The encoder extracts hierarchical features at increasing scales, while the decoder restores the original resolution through upsampling. Skip connections concatenate encoder features with decoder activations, allowing the network to combine high-level context with fine-grained localization details required for precise boundary detection.

Building the U-Net Model in PyTorch

The implementation in lessons/4-ComputerVision/12-Segmentation/SemanticSegmentationPytorch.ipynb defines a UNet class inheriting from nn.Module, utilizing torchvision.models.resnet34 as the backbone.

Encoder Implementation with ResNet34

The encoder leverages a pretrained ResNet34 to extract multi-scale features. The model strips the final fully connected layers and uses the convolutional blocks as feature extractors:

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models

class UNet(nn.Module):
    def __init__(self, n_classes=1):
        super(UNet, self).__init__()
        # Encoder: use pretrained ResNet34

        self.encoder = models.resnet34(pretrained=True)
        self.encoder_layers = list(self.encoder.children())
        self.enc1 = nn.Sequential(*self.encoder_layers[:3])   # 64 channels

        self.enc2 = self.encoder_layers[3]                    # 128 channels

        self.enc3 = self.encoder_layers[4]                    # 256 channels

        self.enc4 = self.encoder_layers[5]                    # 512 channels

        self.enc5 = self.encoder_layers[6]                    # 512 channels (bottleneck)

Decoder with Skip Connections

The decoder uses ConvTranspose2d layers for upsampling and concatenates skip connections via torch.cat. Each decoder block doubles the spatial resolution while halving the channels, then merges with corresponding encoder features:

        # Decoder with upsampling and skip connections

        self.up5 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.dec5 = nn.Conv2d(512, 256, kernel_size=3, padding=1)
        self.up4 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.dec4 = nn.Conv2d(256, 128, kernel_size=3, padding=1)
        self.up3 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.dec3 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
        self.up2 = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
        self.dec2 = nn.Conv2d(64, 32, kernel_size=3, padding=1)
        self.final = nn.Conv2d(32, n_classes, kernel_size=1)

Forward Pass Implementation

The forward method implements the complete data flow through encoder bottleneck and decoder reconstruction:

    def forward(self, x):
        # Encoder path

        e1 = self.enc1(x)          # [B,64,H/2,W/2]

        e2 = self.enc2(e1)         # [B,128,H/4,W/4]

        e3 = self.enc3(e2)         # [B,256,H/8,W/8]

        e4 = self.enc4(e3)         # [B,512,H/16,W/16]

        e5 = self.enc5(e4)         # [B,512,H/32,W/32]

        
        # Decoder path with skip connections

        d5 = self.up5(e5)
        d5 = torch.cat([d5, e4], dim=1)
        d5 = F.relu(self.dec5(d5))
        
        d4 = self.up4(d5)
        d4 = torch.cat([d4, e3], dim=1)
        d4 = F.relu(self.dec4(d4))
        
        d3 = self.up3(d4)
        d3 = torch.cat([d3, e2], dim=1)
        d3 = F.relu(self.dec3(d3))
        
        d2 = self.up2(d3)
        d2 = torch.cat([d2, e1], dim=1)
        d2 = F.relu(self.dec2(d2))
        
        out = self.final(d2)
        return out

Data Loading for Segmentation Tasks

Semantic segmentation requires paired image and mask data. The repository implements a custom SegDataset class extending torch.utils.data.Dataset to load images and their corresponding pixel-wise labels:

import os
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt

class SegDataset(Dataset):
    def __init__(self, img_dir, mask_dir, transform=None):
        self.img_dir = img_dir
        self.mask_dir = mask_dir
        self.transform = transform
        self.images = sorted(os.listdir(img_dir))

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.images[idx])
        mask_path = os.path.join(self.mask_dir, self.images[idx].replace('.jpg', '_mask.png'))
        image = plt.imread(img_path)
        mask = plt.imread(mask_path)
        if self.transform:
            image = self.transform(image)
            mask = self.transform(mask)
        return image, mask

Training the U-Net Model

Training semantic segmentation models requires pixel-level loss calculation and specialized metrics to evaluate mask quality.

Loss Function and Metrics

The implementation uses Binary Cross Entropy with Logits (nn.BCEWithLogitsLoss) for binary segmentation tasks. For evaluation, the dice_coef function measures overlap between predicted and ground truth masks:

def dice_coef(pred, target, smooth=1.):
    pred = torch.sigmoid(pred)
    pred = pred > 0.5
    pred = pred.float()
    intersection = (pred * target).sum()
    return (2. * intersection + smooth) / (pred.sum() + target.sum() + smooth)

Complete Training Loop

Instantiate the model, define the optimizer, and iterate through the dataloader to update weights via backpropagation:

import torch.optim as optim

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

# Assuming train_loader is initialized with SegDataset

for epoch in range(num_epochs):
    model.train()
    epoch_loss = 0
    for imgs, masks in train_loader:
        imgs, masks = imgs.to(device), masks.to(device)
        
        optimizer.zero_grad()
        outputs = model(imgs)
        loss = criterion(outputs, masks)
        loss.backward()
        optimizer.step()
        
        epoch_loss += loss.item()
    
    print(f'Epoch {epoch+1}, Loss: {epoch_loss/len(train_loader):.4f}')

Summary

• U-Net architecture combines a ResNet34 encoder with a transposed convolution decoder to perform end-to-end semantic segmentation.
• Skip connections using torch.cat preserve spatial details by merging encoder features with decoder activations at matching resolutions.
• Binary Cross Entropy with Logits serves as the standard loss function for binary segmentation tasks in the PyTorch implementation.
• Dice coefficient provides a robust metric for evaluating mask overlap quality beyond simple pixel accuracy.
• The complete working implementation resides in microsoft/AI-For-Beginners/lessons/4-ComputerVision/12-Segmentation/SemanticSegmentationPytorch.ipynb.

Frequently Asked Questions

What is the difference between semantic segmentation and instance segmentation?

Semantic segmentation assigns a class label to every pixel without distinguishing between separate objects of the same class. Instance segmentation extends this by individually labeling each distinct object instance. The U-Net implementation in the AI-For-Beginners repository performs semantic segmentation, predicting a single binary mask for foreground versus background classification.

Why does the implementation use ResNet34 as the encoder backbone?

ResNet34 provides pretrained feature extractors that capture hierarchical visual patterns from low-level edges to high-level object parts. Using a pretrained encoder accelerates convergence and improves accuracy compared to training from random initialization, leveraging transfer learning from ImageNet. The residual connections in ResNet34 also mitigate vanishing gradients when extracting deep features at the bottleneck layer.

How do skip connections improve segmentation accuracy?

Skip connections concatenate high-resolution encoder features with upsampled decoder features, compensating for spatial information lost during max-pooling operations. Without these connections, the decoder would reconstruct masks from heavily compressed bottleneck representations alone, resulting in blurry boundaries. The concatenation operation at each decoder stage ensures fine-grained details from early convolutional layers directly influence the final pixel predictions.

When should I use the Dice coefficient instead of pixel accuracy?

Use the Dice coefficient when dealing with imbalanced datasets where foreground pixels comprise a small fraction of the image, such as medical imaging or defect detection. Pixel accuracy can be misleadingly high when the model predicts mostly background. The Dice coefficient specifically measures the overlap between predicted and target regions, providing a balanced metric that equally weights false positives and false negatives through the intersection calculation.

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