# Semantic Segmentation with U-Net Architecture Implementation: A PyTorch Guide

> Implement semantic segmentation using U-Net in PyTorch with Microsoft's AI-For-Beginners guide. Explore encoder-decoder design and skip connections for precise spatial accuracy.

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

---

**The Microsoft AI-For-Beginners repository provides a complete PyTorch implementation of the U-Net architecture for semantic segmentation, featuring an encoder-decoder design with skip connections that concatenate high-resolution encoder features with upsampled decoder maps to preserve spatial precision.**

The `microsoft/AI-For-Beginners` curriculum introduces **semantic segmentation with U-Net architecture implementation** in lesson 4-ComputerVision/12-Segmentation. This hands-on tutorial demonstrates how to build a fully convolutional network that classifies every pixel in an image, using the hallmark U-Net design that combines deep feature extraction with precise localization through skip connections.

## U-Net Architecture Components

The implementation follows the original U-Net paper’s encoder-decoder structure, with four main components working sequentially to transform an input image into a pixel-wise classification mask.

### Encoder Path for Hierarchical Feature Extraction

The encoder progressively down-samples the input through a series of convolutional blocks. In `SemanticSegmentationPytorch.ipynb`, the `ConvBlock` class encapsulates two rounds of `Conv2d → BatchNorm2d → ReLU` operations:

```python
class ConvBlock(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )
    def forward(self, x):
        return self.conv(x)

```

The encoder in the `UNet` class stacks four of these blocks (outputting 64, 128, 256, and 512 channels respectively), applying `nn.MaxPool2d(2)` between each to halve spatial dimensions. The notebook notes that you can swap this simple encoder for a **pretrained ResNet-50** backbone to leverage transfer learning for richer feature representations.

### Decoder Path with Skip Connections

The decoder mirrors the encoder using `nn.ConvTranspose2d` layers for upsampling. The critical **skip connections** are implemented via `torch.cat`, which concatenates the upsampled decoder features with the corresponding encoder feature maps from the same resolution level:

```python
d3 = self.up3(e4)              # Upsample from bottleneck

d3 = torch.cat([d3, e3], dim=1) # Skip connection: concat with encoder stage 3

d3 = self.dec3(d3)             # Refine with convolution

```

This concatenation preserves fine-grained spatial details that would otherwise be lost during the downsampling process, enabling precise boundary detection in the final segmentation mask.

## Complete U-Net Model Implementation

The full `UNet` class in `lessons/4-ComputerVision/12-Segmentation/SemanticSegmentationPytorch.ipynb` ties these components together:

```python
class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1):
        super().__init__()
        self.enc1 = ConvBlock(in_channels, 64)
        self.enc2 = ConvBlock(64, 128)
        self.enc3 = ConvBlock(128, 256)
        self.enc4 = ConvBlock(256, 512)
        
        self.pool = nn.MaxPool2d(2)
        
        self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.dec3 = ConvBlock(512, 256)
        self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.dec2 = ConvBlock(256, 128)
        self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.dec1 = ConvBlock(128, 64)
        
        self.final = nn.Conv2d(64, out_channels, kernel_size=1)
    
    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool(e1))
        e3 = self.enc3(self.pool(e2))
        e4 = self.enc4(self.pool(e3))
        
        d3 = self.up3(e4)
        d3 = torch.cat([d3, e3], dim=1)
        d3 = self.dec3(d3)
        
        d2 = self.up2(d3)
        d2 = torch.cat([d2, e2], dim=1)
        d2 = self.dec2(d2)
        
        d1 = self.up1(d2)
        d1 = torch.cat([d1, e1], dim=1)
        d1 = self.dec1(d1)
        
        return self.final(d1)

```

The final `1×1` convolution maps the 64-channel decoder output to the desired number of classes (`out_channels=1` for binary segmentation, or higher for multi-class tasks).

## Training Configuration and Loss Functions

The notebook implements a standard PyTorch training loop using **binary cross-entropy** loss for single-class masks or **categorical cross-entropy** for multi-class segmentation. Both losses are computed per-pixel and averaged across the batch:

```python
def train_one_epoch(model, loader, optimizer, criterion, device):
    model.train()
    epoch_loss = 0
    for imgs, masks in loader:
        imgs = imgs.to(device)
        masks = masks.to(device)
        
        optimizer.zero_grad()
        preds = model(imgs)
        loss = criterion(preds, masks)
        loss.backward()
        optimizer.step()
        
        epoch_loss += loss.item()
    return epoch_loss / len(loader)

```

The primary evaluation metric demonstrated is **pixel accuracy** (the proportion of correctly classified pixels). The notebook also references additional metrics such as Intersection over Union (IoU) and Dice coefficient for more rigorous model evaluation.

## Dataset and Preprocessing

The tutorial utilizes the **Segmentation Full Body MADS Dataset** from Kaggle, which provides paired images and human body segmentation masks. The data pipeline includes standard augmentations and normalization to prepare inputs for the network.

## Inference and Visualization

For prediction, the model runs in evaluation mode with gradient computation disabled. The notebook demonstrates applying a sigmoid activation and 0.5 threshold to generate binary masks, then visualizing results alongside ground truth:

```python
model.eval()
with torch.no_grad():
    sample_img, _ = next(iter(val_loader))
    pred_mask = torch.sigmoid(model(sample_img.to(device))) > 0.5
    # Visualization code displays input, ground truth, and prediction

```

## TensorFlow Implementation Alternative

Learners preferring TensorFlow can reference `SemanticSegmentationTF.ipynb` in the same directory (`lessons/4-ComputerVision/12-Segmentation/`), which provides an equivalent implementation using Keras layers and the Functional API.

## Summary

- The `microsoft/AI-For-Beginners` repository provides a complete, executable U-Net implementation in `lessons/4-ComputerVision/12-Segmentation/SemanticSegmentationPytorch.ipynb`.
- The architecture uses an encoder-decoder structure with **skip connections** implemented via `torch.cat` to preserve spatial resolution.
- The `ConvBlock` class provides reusable `Conv2d → BatchNorm → ReLU` stacks, while `ConvTranspose2d` handles upsampling in the decoder.
- Loss functions support both binary and multi-class segmentation, with pixel accuracy as the primary metric.
- Optional **ResNet-50 encoder** can be substituted for the simple convolutional encoder to improve feature extraction capabilities.

## Frequently Asked Questions

### What dataset is used in the AI-For-Beginners U-Net tutorial?

The notebook uses the **Segmentation Full Body MADS Dataset** from Kaggle, which contains images paired with pixel-level body segmentation masks. This dataset provides a clear binary classification task suitable for learning semantic segmentation fundamentals.

### How do skip connections function in this U-Net implementation?

Skip connections bridge the encoder and decoder by concatenating feature maps of matching spatial dimensions. In the `forward` method, encoder outputs (`e1`, `e2`, `e3`) are saved and concatenated with their corresponding decoder stages using `torch.cat([decoder_output, encoder_output], dim=1)`, preserving fine details for precise mask boundaries.

### Can I use a pretrained backbone with this U-Net code?

Yes. While the default implementation uses a simple stack of `ConvBlock` layers, the notebook comments indicate you can replace the encoder with a **pretrained ResNet-50** (or other backbone) to leverage transfer learning. This requires adjusting the channel dimensions in the skip connections to match the backbone’s feature pyramid.

### What is the difference between the PyTorch and TensorFlow implementations?

Both implementations follow identical U-Net architectures with encoder-decoder skip connections. The PyTorch version (`SemanticSegmentationPytorch.ipynb`) uses custom `nn.Module` classes and explicit training loops, while the TensorFlow version (`SemanticSegmentationTF.ipynb`) utilizes the Keras Functional API with built-in `model.fit` training methods. The core architectural concepts remain consistent across both frameworks.