U-Net, nnU-Net, and Transformer-Based Segmentation

Module: Healthcare AI | Difficulty: Advanced

Dice Loss

Focal Tversky Loss

Boundary Loss

where is the signed distance transform of the ground truth and is the predicted segmentation map.

Segmentation Architecture Comparison

| Architecture | Params | Dice Score | Speed | Auto-config | |-------------|--------|------------|-------|-------------| | U-Net | 31M | 0.86 | Fast | No | | Attention U-Net | 35M | 0.88 | Medium | No | | nnU-Net | Variable | 0.91 | Medium | Yes | | Swin-UNet | 27M | 0.89 | Slow | No | | UNETR | 100M+ | 0.90 | Slow | No |

import torch
import torch.nn as nn

class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.conv(x)

class UNet(nn.Module):
    def __init__(self, in_channels=1, num_classes=1):
        super().__init__()
        self.enc1 = DoubleConv(in_channels, 64)
        self.enc2 = DoubleConv(64, 128)
        self.enc3 = DoubleConv(128, 256)
        self.enc4 = DoubleConv(256, 512)
        self.pool = nn.MaxPool2d(2)
        self.bottleneck = DoubleConv(512, 1024)
        self.up4 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
        self.dec4 = DoubleConv(1024, 512)
        self.up3 = nn.ConvTranspose2d(512, 256, 2, stride=2)
        self.dec3 = DoubleConv(512, 256)
        self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.dec2 = DoubleConv(256, 128)
        self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.dec1 = DoubleConv(128, 64)
        self.out_conv = nn.Conv2d(64, num_classes, 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))
        b = self.bottleneck(self.pool(e4))
        d4 = self.dec4(torch.cat([self.up4(b), e4], dim=1))
        d3 = self.dec3(torch.cat([self.up3(d4), e3], dim=1))
        d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1))
        d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
        return self.out_conv(d1)

model = UNet(in_channels=1, num_classes=1)
x = torch.randn(1, 1, 256, 256)
print(f'Output shape: {model(x).shape}')

Research Insight: nnU-Net automatically configures the entire segmentation pipeline based on dataset properties. It meta-learns from a large corpus of medical segmentation benchmarks and adapts its configuration to new datasets without manual hyperparameter tuning, achieving state-of-the-art results across 23 public benchmarks.