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

> Learn to implement semantic segmentation in PyTorch using U-Net models. Build a ResNet34 encoder, transposed convolutional decoder, and skip connections. Train efficiently with binary cross-entropy loss.

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

---

**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:

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

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

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

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

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

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