Model Architecture

Model Architecture

The primary model is a 3D Attention U-Net with squeeze-and-excitation channel attention and spatial attention gates. A lighter 2D U-Net is also included for quick experiments.

---

3D Attention U-Net (src/segmentation/unet3d.py)

from src.segmentation import AttentionUNet3D

model = AttentionUNet3D(
    in_channels=1,
    num_classes=8,
    base_filters=32,
    depth=2,
    dropout=0.1
)
# Input:  (B, 1, D, H, W)
# Output: (B, 8, D, H, W)
# Classes:
#   0 = background
#   1 = Glioma            (Astrocytoma, Glioblastoma, Oligodendroglioma)
#   2 = Meningioma        (Meningothelial Tumors)
#   3 = Nerve Sheath      (Schwannoma, Neurocytoma)
#   4 = Embryonic         (Medulloblastoma, DNET)
#   5 = Mixed Neuronal    (Ependymoma, Ganglioglioma)
#   6 = Mesenchymal       (Hemangiopericytoma)
#   7 = Germ Cell         (Germinoma)

Full Architecture Diagram (depth=2, base_filters=32 — the trained CPU config)

flowchart TB
    A0["Input Volume<br/>(B, 1, D, H, W)<br/>e.g. (B, 1, 2, 64, 64)"]

    A0 --> E0["Encoder L0<br/>Conv3DBlock 1→32 + SE Attention"]
    E0 -- "skip 0 (32ch, full res)" --> D1
    E0 --> P1["MaxPool3D (1,2,2)"]

    P1 --> E1["Encoder L1<br/>Conv3DBlock 32→64 + SE Attention"]
    E1 -- "skip 1 (64ch, 1/2 res)" --> D0
    E1 --> P2["MaxPool3D (1,2,2)"]

    P2 --> E2["Encoder L2<br/>Conv3DBlock 64→128 + SE Attention"]
    E2 --> Bn["Bottleneck<br/>Conv3DBlock 128→128 + Dropout"]

    Bn --> U0["ConvTranspose3D (1,2,2)<br/>128→64"]
    U0 --> D0["Attention Gate + Concat (128ch)<br/>Conv3DBlock 128→64"]

    D0 --> U1["ConvTranspose3D (1,2,2)<br/>64→32"]
    U1 --> D1["Attention Gate + Concat (64ch)<br/>Conv3DBlock 64→32"]

    D1 --> F["Final Conv3D 1×1×1<br/>32 → 8 classes"]
    F --> O["Output Logits<br/>(B, 8, D, H, W) → argmax → mask"]

    classDef enc fill:#3b82f6,color:#fff,stroke:none
    classDef dec fill:#10b981,color:#fff,stroke:none
    classDef bot fill:#f59e0b,color:#fff,stroke:none
    classDef io fill:#6b7280,color:#fff,stroke:none
    class E0,E1,E2,P1,P2 enc
    class U0,U1,D0,D1 dec
    class Bn bot
    class A0,F,O io

Loading

With depth=N in general, the pattern repeats: N encoder levels (each halving H/W, doubling channels), a bottleneck, then N decoder levels (each undoing one encoder level via a matching skip connection). depth=2 (CPU preset) gives 2 pooling stages; the 8GB/16GB GPU presets use depth=3/4 for a deeper receptive field.

Encoder

Each encoder level consists of:

1. Conv3DBlock — two 3×3×3 convolutions with BatchNorm + ReLU
2. Squeeze-and-Excitation (SE) block — channel attention: global average pool → FC → sigmoid → scale
3. Down3D — MaxPool3d(kernel_size=(1,2,2)) to downsample H and W only, preserving depth

The (1,2,2) kernel is intentional. The model processes pseudo-3D inputs with small depth (D=2), so standard (2,2,2) pooling would collapse the depth dimension by the second encoder level.

Conv3DBlock internals

flowchart LR
    x["Input<br/>C_in channels"] --> c1["Conv3d 3×3×3"] --> bn1["BatchNorm3D"] --> r1["ReLU"]
    r1 -.-> drop["Dropout3D<br/>(bottleneck only)"]
    r1 --> c2["Conv3d 3×3×3"]
    drop --> c2
    c2 --> bn2["BatchNorm3D"] --> se["SE Channel<br/>Attention"]
    se --> add(("+"))
    x -- "residual<br/>(1×1 conv if channels differ)" --> add
    add --> r2["ReLU"] --> y["Output<br/>C_out channels"]

Loading

SE Channel Attention internals

flowchart LR
    x["Feature Map<br/>(C, D, H, W)"] --> gap["Global Avg Pool3D<br/>→ (C, 1, 1, 1)"]
    gap --> fc1["Conv3d 1×1×1<br/>C → C/16"] --> relu["ReLU"] --> fc2["Conv3d 1×1×1<br/>C/16 → C"] --> sig["Sigmoid"]
    sig --> scale(("×"))
    x --> scale
    scale --> y["Recalibrated<br/>Feature Map"]

Loading

Bottleneck

Conv3DBlock at the deepest level with higher dropout.

Decoder

Each decoder level:

1. Up3D — ConvTranspose3d(kernel_size=(1,2,2)) to upsample H and W only
2. Attention gate on the skip connection — suppresses irrelevant spatial locations before concatenation
3. Conv3DBlock — fuse upsampled features + attended skip connection

Attention Gate internals

flowchart LR
    g["Gating Signal g<br/>(upsampled decoder features)"] --> wg["Conv3d 1×1×1 + BN"]
    xskip["Skip Connection x<br/>(from encoder)"] --> wx["Conv3d 1×1×1 + BN"]
    wg --> add(("+"))
    wx --> add
    add --> relu["ReLU"] --> psi["Conv3d 1×1×1<br/>+ BN + Sigmoid"]
    psi --> mult(("×"))
    xskip --> mult
    mult --> y["Attended Skip<br/>Connection"]

Loading

The gate learns a per-location weight in [0, 1] from the combination of the decoder's gating signal and the encoder's skip features, then scales the skip connection before it's concatenated — this suppresses background/irrelevant regions the encoder would otherwise pass straight through.

Output

1×1×1 convolution → 8-channel logits → argmax at inference time.

---

Pseudo-3D Training

The Kaggle dataset contains 2D MRI slices, not full 3D volumes. To use the 3D model, each slice is stacked D_FRAMES times to create a (B, 1, D, H, W) volume:

D_FRAMES = 2  # configurable via --d-frames or config JSON
images_3d = images.unsqueeze(2).repeat(1, 1, D_FRAMES, 1, 1)

This is a pragmatic bridge — the model learns spatial context within a slice at multiple scales. With true volumetric data (e.g., BraTS), D_FRAMES would be replaced by real depth.

---

2D U-Net (src/segmentation/unet.py)

from src.segmentation import UNet

model = UNet(in_channels=1, n_classes=1, base_filters=32, depth=4)
# Input:  (B, 1, H, W)
# Output: (B, 1, H, W)  — binary logits

Used for quick single-class training (scripts/train.py). Binary output — tumor vs background.

---

Parameter Count

Model	Depth	Base Filters	Parameters
2D U-Net	4	32	~7M
3D Attention U-Net (CPU preset)	2	32	2.2M
3D Attention U-Net (8GB GPU preset)	3	32	~15M
3D Attention U-Net (16GB GPU preset)	4	32	~31M

---

Loss Functions (src/segmentation/losses_advanced.py)

Loss	Purpose
WeightedDiceLoss	Per-class Dice with class weights — handles imbalance
FocalLoss	Down-weights easy background pixels
BoundaryLoss	Penalizes imprecise tumor edges
HybridLoss	α·WeightedDice + β·Focal + γ·Boundary

Default weights in train3d.py: α=0.5, β=0.3, γ=0.2.

Class weights for 8-class: [0.01, 0.18, 0.18, 0.15, 0.14, 0.14, 0.11, 0.09] (background gets near-zero weight; rarer classes like Germ Cell and Mesenchymal get higher weights).