পাঠ ২২ · ৩৫-এর মধ্যে · মডিউল ৩

Home / AI Courses / Computer Vision / U-Net

U-Net — segmentation-এর জাদু

U-Net — encoder-decoder with skip connections

৮ মিনিট পড়া মাঝারি · Intermediate PyTorch কোডসহ

এই পাঠে যা শিখবেন

U-Net architecture — encoder, decoder, skip
Why skip connection matters
U-Net variants — Attention U-Net, U-Net++, nnU-Net
PyTorch implementation

১ · U-Net-এর জন্ম

Olaf Ronneberger, Philipp Fischer, Thomas Brox — Freiburg University, ২০১৫। ISBI cell tracking challenge winner. MICCAI conference paper.

Original problem: EM (electron microscope) cell segmentation. Only 30 training image. Need precise boundary.

কেন্দ্রীয় ধারণা

U-Net = encoder (downsample, abstract context) + decoder (upsample, precise localization) + skip connection (pass detail across). U-shape architecture.

২ · U-Net architecture

Contracting path (encoder)

4 stage — each: Conv 3×3 → ReLU → Conv 3×3 → ReLU → MaxPool 2×2.
Channel double per stage: 64 → 128 → 256 → 512 → 1024.
Spatial halve per stage.

Expansive path (decoder)

4 stage — each: TransposeConv 2×2 (upsample) → concat skip → Conv 3×3 → ReLU → Conv 3×3 → ReLU.
Channel halve per stage: 1024 → 512 → 256 → 128 → 64.
Spatial double per stage.

Skip connections

Encoder feature → decoder same-resolution stage.
Concatenate (channel-wise) — not add.
Feature map pass — detail preserve.

Output

1×1 conv → number of classes (typically 1 for binary, 2 for foreground/background).
Original 388×388 (slightly smaller than input 572×572 due to "valid" conv).
Modern padding-same — same size output.

৩ · কেন U-Net medical-এ revolutionary

Small dataset: 30 image-এ train possible — extreme augmentation.
Boundary precision: skip connection — pixel-level accuracy.
Class imbalance: weighted loss handle.
Tile-based: large image tile-এ process.
Generic: any modality (CT, MRI, X-ray) adapt.

৪ · Skip connection-এর math

ResNet-এর add নয় — concatenation:

$$\text{decoder}_l = \text{Conv}\big( \text{concat}(\text{upsample}(\text{decoder}_{l-1}), \text{encoder}_l) \big)$$

Encoder $l$ — high resolution, low semantic.
Decoder upsampled — low resolution origin, semantic enriched.
Concat — both signal preserve.
Subsequent conv — combine.

U-Net = "guide a sketch artist who's learnt big picture but lost fine detail". Encoder gives understanding (what is this?), decoder reconstructs precise boundary, skip connection passes back fine line.

U-Net architecture — contracting path-এ context, expansive path-এ localization, skip connection-এ detail.

৫ · PyTorch implementation

Python · PyTorch

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

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, n_classes=1):
        super().__init__()
        self.enc1 = DoubleConv(3, 64)
        self.enc2 = DoubleConv(64, 128)
        self.enc3 = DoubleConv(128, 256)
        self.enc4 = DoubleConv(256, 512)
        self.bottleneck = DoubleConv(512, 1024)
        self.up4 = nn.ConvTranspose2d(1024, 512, 2, 2)
        self.dec4 = DoubleConv(1024, 512)
        self.up3 = nn.ConvTranspose2d(512, 256, 2, 2)
        self.dec3 = DoubleConv(512, 256)
        self.up2 = nn.ConvTranspose2d(256, 128, 2, 2)
        self.dec2 = DoubleConv(256, 128)
        self.up1 = nn.ConvTranspose2d(128, 64, 2, 2)
        self.dec1 = DoubleConv(128, 64)
        self.out = nn.Conv2d(64, n_classes, 1)
    def forward(self, x):
        e1 = self.enc1(x);              p1 = F.max_pool2d(e1, 2)
        e2 = self.enc2(p1);             p2 = F.max_pool2d(e2, 2)
        e3 = self.enc3(p2);             p3 = F.max_pool2d(e3, 2)
        e4 = self.enc4(p3);             p4 = F.max_pool2d(e4, 2)
        b = self.bottleneck(p4)
        d4 = self.dec4(torch.cat([self.up4(b), e4], 1))
        d3 = self.dec3(torch.cat([self.up3(d4), e3], 1))
        d2 = self.dec2(torch.cat([self.up2(d3), e2], 1))
        d1 = self.dec1(torch.cat([self.up1(d2), e1], 1))
        return self.out(d1)

m = UNet(n_classes=2)
x = torch.randn(1, 3, 256, 256)
print("Output:", m(x).shape)         # (1, 2, 256, 256)
print(f"Params: {sum(p.numel() for p in m.parameters()):,}")

~31M parameters. Output spatial সমান input. Per-pixel logits → softmax → class probabilities.

৬ · U-Net variants

Attention U-Net: skip connection-এ attention gate.
U-Net++: nested skip — denser connection.
3D U-Net: volumetric medical (CT, MRI).
V-Net: 3D, residual connection.
nnU-Net (২০২০): "no new U-Net" — automated configuration. Medical SOTA.
TransUNet: Transformer encoder + U-Net decoder.
SwinU-Net: Swin Transformer-based.

৭ · Stable Diffusion-এ U-Net

Image generation-এ U-Net key role. Stable Diffusion (২০২২)-এর core network — U-Net (with cross-attention).

Noise prediction U-Net — denoising step iteratively.
Cross-attention layer — text condition inject.
Skip connection — generation detail preserve.

৮ · Medical applications

Tumor segmentation: brain (BraTS), liver, breast.
Organ segmentation: liver, kidney, heart.
Cell counting: microscopy.
Retinal vessel: diabetic retinopathy detect.
Bone fracture: X-ray.
Chest X-ray: tuberculosis, COVID lesion.
Bangladesh: JIBON Bangladesh, ICDDR,B এ U-Net deploy হচ্ছে।

U-Net medical-এ default — কিন্তু input-output size identical চাই। Padding strategy ('same' padding modern) ও tile-based inference for large image.

ভাবনার প্রশ্ন

প্র ০১ U-Net-এর skip connection-এ concat (not add) কেন? ResNet-এর add কেন এখানে inappropriate?

এটি architectural choice-এর key insight.

ResNet add:

$y = F(x) + x$ — channel count same require.
"Residual" learn — small refinement to identity.
Information lossy if F summarizes weakly.

U-Net concat:

$y = \text{conv}(\text{cat}(F(x), x))$ — channel double.
Both signal preserved.
Subsequent conv combines both.
"Encoder feature + decoder feature" — different role.

Why concat better for U-Net:

Encoder feature: high-res detail, low semantic.
Decoder feature: low-res, high semantic (post-bottleneck).
These are different in nature — sum-এ blend lose info.
Concat preserve both, network combines optimally.

Trade-offs:

Concat: more parameters (channel double).
Add: parameter-equal, potentially less expressive.

Empirical:

U-Net add — 1-2% mIoU drop.
Original paper concat চuose — Ronneberger evidence.

Modern alternative:

FPN — top-down + lateral. Add or concat both used.
BiFPN — weighted concat-with-attention.

মূল উপলব্ধি: "Skip connection" — generic concept, implementation matters. Use case-specific choice.

প্র ০২ Original U-Net trained on only 30 images! এত small data-এ deep network train possible — কী trick?

এটি U-Net-এর অসাধারণ achievement. Few-shot deep learning-এর precursor.

Key tricks:

Aggressive augmentation: elastic deformation, rotation, shift, flip. One image → 100s effective.
Architecture inductive bias: CNN's translation equivariance — implicit regularization.
Tile-based training: large image → many patches. Each patch independent example.
Weighted loss: rare class boost.

Elastic deformation specifically:

Random displacement field per pixel.
Smoother than affine — natural-looking distortion.
Mimics tissue elasticity in microscopy.
Critical for medical small-data regime.

Why architecture helps:

Skip connection → gradient flow easy. Deep network train possible.
Encoder-decoder — "compose" representation.
Per-pixel output — many supervision signal per image.

Modern small-data strategies:

Pretrained encoder: ImageNet ResNet — transfer.
Self-supervised pretrain: domain-specific MAE.
Few-shot meta-learning: MAML, ProtoNet.
SAM: foundation segmentation-এ from clicks.

Bangladesh medical data:

Most hospital dataset 100-1000 image.
U-Net + augment + transfer = robust baseline.
Federated learning — multi-hospital privacy-preserving.

Numbers from paper:

30 train image, 30 test.
EM cell IoU: 0.93.
2015-এ revolutionary.

মূল উপলব্ধি: "Big data" myth challenge. Architecture-aware augmentation small-data DL সম্ভব। Bangladesh medical AI — U-Net foundation.

প্র ০৩ nnU-Net (২০২০) — "no new U-Net" — automatic configuration. কী magic? Manual tuning থেকে কী gain?

nnU-Net (Isensee et al., Nature Methods, 2021) — medical segmentation revolutionary. Manual hyperparameter tuning automate.

The problem:

Each medical dataset different: modality (CT vs MRI), spacing, intensity range, class.
Per-dataset hyperparameter tune — extensive expertise.
Inconsistent baseline across paper.

nnU-Net solution:

Dataset analyze automatically.
Optimal architecture choose: 2D vs 3D, patch size, batch size.
Optimal preprocessing: resampling, normalization.
Optimal training: optimizer, LR, augmentation.
Self-configuring pipeline.

Key components:

Fingerprint: dataset statistics extract.
Rule-based config: heuristic from fingerprint.
Empirical evidence: 53 dataset paper basis.
Cross-validation: 5-fold standard.
Ensemble: 5 fold model average.

Performance:

10+ MICCAI challenge winner.
Often outperforms SOTA bespoke architecture.
"Don't reinvent — properly configure."

Implication:

Architecture innovation overrated, configuration underrated.
Hand-tuning expert-time intensive.
Reproducibility — single tool, many dataset.

Adoption:

nnU-Net — medical research standard.
Open source, PyTorch.
NIH benchmark default.

Bangladesh use:

Hospital research lab — nnU-Net plug-and-play.
BIRDEM, ICDDR,B, BSMMU — nnU-Net pilots.

Lesson for engineers:

Automation > novelty.
Empirical evidence > anecdote.
Reproducibility critical.

মূল উপলব্ধি: "ML system" > "ML model". nnU-Net = configuration as code, validated at scale. Modern AI engineering principle.

প্র ০৪ Stable Diffusion — image generation — কেন U-Net? Discriminative segmentation থেকে generative-এ U-Net কীভাবে adapt?

U-Net-এর pivot — Rombach et al. (২০২২) Stable Diffusion paper-এ। Surprising adoption.

Diffusion model basics:

Forward process: image-এ gradually noise add.
Reverse process: noise থেকে gradually image generate.
Need: noise prediction network for each step.

U-Net role:

Input: noisy image (or latent in SD).
Output: predicted noise.
Same input-output size — U-Net structure ideal.

Why U-Net specifically?

Same size in/out: diffusion need this.
Skip connection: fine detail preserve during denoising.
Hierarchy: noise at multiple scale.
Available code: pretrained, well-understood.

Modifications for SD:

Attention layer: self-attention at each scale (transformer-like).
Cross-attention: text embedding inject — language condition.
Time embedding: denoising step encode.
Latent space: SD operates 4×64×64 latent (not pixel).

SD U-Net specifics:

~860M params (SD 1.5).
4 down + 4 up stage.
Text condition via cross-attention.
Trained on LAION-5B images.

Modern evolution:

SD 3 (2024): DiT (Diffusion Transformer) — pure transformer, U-Net out.
FLUX: hybrid.
U-Net being phased out for transformer in image generation.

U-Net legacy:

৭ year (২০১৫-২০২২) medical dominance.
Generative AI bridge era — Stable Diffusion 1, 2.
Now niche but historically pivotal.

Bangladesh implication:

Stable Diffusion fine-tune (Bangla art).
U-Net + ControlNet — controllable generation.
LoRA — task-specific adaptation.

মূল উপলব্ধি: Architecture transferable across paradigm — discriminative → generative. U-Net's geometric flexibility makes it adaptable. Engineering principle: simple shapes (U) compose well.

অনুশীলন

Param count: Original U-Net (3 → 64 → 128 → 256 → 512 → 1024) total parameter approximate.

~31M parameters. Mostly bottleneck (512×1024 conv 3×3) ও দু'টি conv per stage.
Modify: U-Net-এর encoder ResNet-50 দিয়ে replace করার সংক্ষিপ্ত description.

ResNet-50 stage 1-4 use as encoder, skip connection at each stage. Decoder same. segmentation_models_pytorch library এই pre-built.
ভাবুন: Bangladesh chest X-ray TB detect — U-Net কেন appropriate? Class imbalance handle কীভাবে?

U-Net medical-এ pretrained transfer + small-data robust. TB lesion 2-5% pixel — Dice + weighted CE loss combine. Transfer learning ImageNet ResNet encoder.