phd-nnunet

PhD Research Progress Summary

Medical Image Segmentation on Intel Gaudi HPUs with nnUNet

Last Updated: 2026-05-03
Dataset: ACDC (Dataset927_ACDC) — Cardiac MRI segmentation
Classes: RV (Right Ventricle), MLV (Myocardium), LVC (Left Ventricle Cavity)
Base Model: nnUNet v2 ResEncUNetXL (3D full resolution, fold 0)

---

Session 1: Data Preparation & Repository Setup

What Was Done

• Established three benchmark datasets: KiTS2023 (kidney), AMOS2022 (multi-organ), ACDC (cardiac)
• Set up dual framework comparison: nnUNet v2 vs Advanced_nnUNet
• Created result analysis tooling (summarize_nnunet_results.py — 750+ line Excel pivot table generator)

Critical Issues Found & Fixed

Key Conclusions

• Dataset927 (ACDC) is the only dataset with truly identical data + splits across both frameworks → only valid apples-to-apples comparison
• Dataset920 (KiTS2023) has same subset but nnUnet has 25% more data → comparison valid with caveat
• Dataset923 (AMOS) has different class counts (13 vs 15) → not directly comparable

---

Session 2: Architecture Changes (50-Epoch Benchmarks)

Experimental Setup

• Protocol: 50-epoch screening runs on Dataset927_ACDC, fold 0, 3D full resolution
• Plan: nnUNetResEncUNetXLPlans
• Goal: Identify best architectural modifications before full 1000-epoch training

Completed Results (Ranked by Mean Dice)

Still Running (as of 2026-03-14 morning)

Crashed at Epoch 0 (need bug fixes)

Failed Plan Variants (data incompatibility with ACDC)

Key Observations

1. LabelSmooth is the clear winner (+1.0% over baseline) — simple regularization very effective
2. FullKernel (larger convolution kernels) and DiceHeavy loss close behind
3. FatDecoder (wider decoder) and ResDenseUNet (hybrid architecture) show architecture changes work
4. AdamW optimizer consistently helps across loss variants
5. All successful trainers cluster in 91.9%–92.9% — ceiling effect at ~93%
6. 7 architectural experiments crashed and need code fixes before re-running
7. Plan variants fail catastrophically on ACDC — unnecessary for this dataset
8. Performance ceiling estimated at ~0.93 Dice for fold 0, 50-epoch screening

---

Session 3: Inference Optimization

Baseline

Model: nnUNetTrainer__nnUNetResEncUNetXLPlans__3d_fullres, fold 0
Default inference Dice: 0.9111 (checkpoint_final)

Round 1: Basic Inference Modifications (7 experiments)

Round 2: Advanced Single-Model Techniques (8 experiments)

Baseline: checkpoint_best = 0.9126

Round 3: Multi-Fold Ensemble (2 experiments)

Round 4: Targeted Post-Processing (12 experiments)

Baseline: fold 0 best + anatomy = 0.9139

Round 4b: Step Size Optimization

Result: All step sizes (25%, default, 125%, 375%) produce identical Dice (0.9143) — ACDC volumes too small for step size to matter.

Best Inference Configurations

Key Inference Lessons

1. nnUNet inference is already near-optimal — most modifications yield 0.00% change
2. Checkpoint selection matters (+0.15%) — minor late-epoch overfitting exists
3. Only domain-specific postprocessing helps — anatomical constraints (+0.13%)
4. Multi-fold ensemble is the biggest win (+0.40% over default, +0.18% over tuned single)
5. Generic techniques universally fail: temperature scaling, intensity TTA, probability smoothing all ineffective
6. Persistent failure cases (patient049, patient034, patient091) are model-inherent — need architecture fixes
7. Cascading without training-time support destroys performance (-36.5%)

---

Cross-Framework Comparison (ACDC — Dataset927)

Takeaway: Advanced_nnUNet models outperform nnUNet v2 by +0.58% mean Dice, primarily from better RV segmentation. Attention/hybrid architectures handle the irregular RV shape better.

---

Overall Best Results Achieved

---

Infrastructure Notes

Pod Management

• Always create a new pod for each new experiment — reusing pods across different experiments leads to evictions and wasted GPU-hours
• Same pod OK for: continuing/resuming the same experiment, re-running validation on same model
• New pod needed for: different trainer, different dataset, different architecture
• Pod template: /home/kkaczor/Desktop/scripts/moje/yamls/124_latest_2100_phd.yaml
• Create: hlctl create containers --flavor g3 --file <yaml> --name kkaczor-<task> --namespace framework --username kkaczor -q --no-pager --as-job --shm 10240
• Fresh pods need: pip install -e /software/users/kkaczor/phd/nnUnet/nnUNet
• Data: shared weka storage at /software/users/kkaczor/phd/nnUnet/

Environment Variables (set before training)

export nnUNet_raw="/software/users/kkaczor/phd/nnUnet/nnUNet_raw"
export nnUNet_preprocessed="/software/users/kkaczor/phd/nnUnet/preprocessed"
export nnUNet_results="/software/users/kkaczor/phd/nnUnet/results"

---

Session 4: Ablation Study — Feature Combinations (March 14, 2026)

Design Rationale

From the 50-epoch benchmarks, 4 independent successful features were identified on the ResEncUNetXL architecture, plus ResDenseUNet as a standalone:

All combinations use consistent base: AdamW 1e-3 + CosineAnnealingLR + TopK(k=10) loss.

For the ablation-in-pairs approach (AdamW variants 6-9), AdamW_DiceHeavy (0.9262) was the best; since DiceHeavy is already a separate feature (DH), the AdamW optimizer is standardized as the base.

Full Experiment Matrix

Main ablation path (for publication):

Baseline (0.9194) → +LS (0.9294) → +LS+FK → +LS+FK+DH → +LS+FK+DH+FD

All 14 new experiments:

Trainer Files

All located in: nnunetv2/training/nnUNetTrainer/variants/combinations/

Launch

bash /home/kkaczor/software/phd/nnUnet/launch_ablation_study.sh

Creates 14 pods (one per experiment), each running 50 epochs on fold 0.

---

Session 5: 1000-Epoch Full Training & Combo Ablation Results (May 3, 2026)

Overview

The 50-epoch screening phase is essentially complete with 89 experiments finished. The focus has shifted to full 1000-epoch training of the most promising combinations, with 14 experiments promoted from screening.

Experiment Scale

Top 20 Experiments by Validation Dice (50-epoch, fold 0)

50-epoch Baseline (ResEncUNetXL): 0.9194

Baseline Full Training Results (1000 epochs, 5-fold CV, XL-plan)

Completed 1000-Epoch Experiments (Test Inference)

Note: DeepDec and ComboLS_LateZ_DH used L-plan (smaller model), not XL. Comparison with XL baseline is not apples-to-apples. The XL-plan combo variants are the main challengers.

1000-Epoch Training Status (12 XL-plan experiments — STOPPED, now resuming)

Pod kkaczor-1k-resume-tfjob launched May 3, 2026 on Gaudi 3 (g3) to resume the top 5 experiments.

Training Times

Optimizer Configuration Investigation

All 1000-epoch trainers were investigated for a suspected LR schedule bug (T_max targeting 50 instead of 1000):

• Source code: All 14 trainers override configure_optimizers() with T_max=self.num_epochs (1000) ✅
• debug.json: All show num_epochs=1000 ✅
• Training logs: LR at epoch 50 = 0.00099 (correct for T_max=1000). If buggy, it would be ~1e-6 ✅
• Optimizer: torch.optim.AdamW (switched from FusedAdamW to fix step-counter serialization bug on checkpoint resume)
• Settings: initial_lr=1e-3, weight_decay=3e-2, CosineAnnealingLR, eta_min=1e-6

Verdict: LR schedule fix was applied BEFORE any 1000-epoch training started. No experiments affected by the bug.

What Worked (Key Insights)

1. Label Smoothing (LS) is the single most impactful regularization technique (+1.0% over baseline at 50ep)
2. LateZ kernels (late application of z-axis convolutions) combine very well with LS
3. Deep Decoder with Xavier initialization is the top combination (0.9329 at 50ep)
4. DiceHeavy loss (2:0.5 Dice:CE ratio) consistently helps across all combinations
5. Combinations outperform singles — the ablation path Baseline→+LS→+LateZ→+DH→+DeepDec→+Xavier shows monotonic improvement
6. Reinit strategies (reinitializing weights at epoch 50) work nearly as well as Xavier from scratch
7. AdamW + CosineAnnealingLR is the best optimizer setup (vs SGD+PolyLR baseline)
8. Model size matters at test time — L-plan 1000ep experiments (DeepDec: 0.9043, ComboLS_LateZ_DH: 0.8793) significantly underperform XL baseline (0.9111)

What Didn't Work

1. FatDecoder (FD) — ranked lower in combos than expected; overhead doesn't justify gains
2. Alternative architectures — DenseUNet (0.9034) and TrueDenseResidual (0.9084) trail baseline
3. Plan variants (FineSpacing, LowOrder, NNLabel) — catastrophic failure on ACDC
4. Generic inference postprocessing — temperature scaling, TTA, smoothing all ineffective
5. FocalTversky loss — produced all-zero predictions
6. GELU activation — marginal improvement, not worth the overhead
7. TopK5 loss — too aggressive, slightly hurts
8. L-plan models for 1000ep — insufficient capacity; the 3.18% gap of ComboLS_LateZ_DH shows smaller models can't match XL even with better training tricks

Generalization Status

Cross-dataset experiments on Dataset920 (KiTS2023) and Dataset923 (AMOS2022) have partial results. Some crashed and need reruns. This is tracked in launch_all_missing.py.