feat(benchmark): extend training to 20 epochs and improve checkpoint handling for nnU-Net

BENCHMARKING.md

@@ -176,6 +176,46 @@ test dice 0.1424 (29/39 rows > 0.001, max 0.675). Checkpoints in
 `/home/xfr/longitudinal_nnu/results/Dataset221_T1C_BENCH/...fold_0/`
 (copy to `nnu/results/` for the record if local space is reclaimed).
 
+## Extended to 20 epochs (2026-10-05/06)
+
+All three continued from their 10-epoch checkpoints to a 20-epoch total
+(`--train-epochs 20` = total budget; 10-epoch JSONs snapshotted as
+`results/benchmark_full_{A,B,C}_10ep.json`):
+
+```
+python scripts/benchmark_pipelines.py --pipelines A --gpu 0 --resume --train-epochs 20 \
+    --a-fg-ce-weight 500 --out results/benchmark_full_A.json
+python scripts/benchmark_pipelines.py --pipelines C --gpu 2 --resume --train-epochs 20 \
+    --out results/benchmark_full_C.json
+NNU_ROOT=/home/xfr/longitudinal_nnu BENCH_B_DIR=/home/xfr/longitudinal_benchB \
+python scripts/benchmark_pipelines.py --pipelines B --gpu 1 --resume --nnu-skip-prep \
+    --train-epochs 20 --out results/benchmark_full_B.json
+```
+
+A/C resume from `runs_bench/{A,C}/best.pt` (model + optimizer + scheduler,
+epochs 11–20). B: the completed run had deleted `checkpoint_latest.pth`
+(see Resuming training), so it was restored from `checkpoint_final.pth`
+(`current_epoch: 10`) before `--c` + `--nnu-skip-prep` — optimizer state and
+LR schedule restored (initial validation 0.1253). `bench_B`'s resume
+s/epoch now divides by epochs actually trained (reads `current_epoch` from
+the checkpoint).
+
+| pipeline | dice @10ep | dice @20ep | Δ | median | max | rows > 0.001 |
+|---|---|---|---|---|---|---|
+| B nnU-Net 3d_fullres | 0.1424 | **0.2927** (hard 0.2927) | +0.150 | 0.289 | 0.839 | 29/39 |
+| C MONAI UNet(16→128) | 0.2162 | **0.1621** | −0.054 | 0.126 | 0.680 | 31/39 |
+| A Unet3D improved (fg crops + CE w=500) | 0.0369 | **0.0076** | −0.029 | 0.000 | 0.196 | 6/39 |
+
+Resume window: A 229.2 s/epoch, C 253.9 s/epoch, B 280.1 s/epoch
+(10 epochs each, same data budget). B was far from convergence at 10 epochs
+and improved a lot; C's loss was flat (0.184) and its extra epochs slightly
+degraded test dice; the A w=500 config oscillates (loss 0.49–0.55, pmax
+median 0.55 → 0.49, rows >0.5: 29 → 19), so its 10-epoch point is its best
+under this config. At 20 epochs **B leads (0.2927)**. Checkpoints in
+`runs_bench/{A,C}/best.pt` and B's fold_0 now hold the 20-epoch models — the
+10-epoch models were overwritten (JSONs preserved). B: 21.9 s/vol TTA,
+9.8 GB train VRAM.
+
 ## Resuming training (`--resume`)
 
 All three pipelines can continue from their latest benchmark checkpoint:
@@ -198,7 +238,9 @@ python scripts/benchmark_pipelines.py --pipelines "A B C" --resume \
   preprocessed tree isn't `--clean`ed; the isolated Dataset221 is rebuilt
   otherwise and a finished run's `checkpoint_latest.pth` is gone, so a
   completed run + a bigger budget retrains B from scratch (interrupted runs
-  keep `checkpoint_latest.pth` and resume correctly).
+  keep `checkpoint_latest.pth` and resume correctly). A finished run CAN
+  still be resumed by copying `checkpoint_final.pth` → `checkpoint_latest.pth`
+  first (same state dict; used for the 20-epoch extension above).
 - Verified 2026-10-04: A/C resumed epoch 1 → 3 (model restored, fresh
   opt/sched for the old-format ckpt), re-run at 3/3 trained 0 epochs, B
   continued from epoch 1 with LR restored mid-decay (0.00536 @ start of

scripts/benchmark_pipelines.py

@@ -490,6 +490,12 @@ def bench_B(args, train_rows, test_rows):
             # (= --train-epochs, the total budget)
             bcmd.append("--c")
             res["resumed_from_checkpoint_latest"] = True
+            try:
+                _st = torch.load(os.path.join(N.fold_dir(), "checkpoint_latest.pth"),
+                                 map_location="cpu", weights_only=False)
+                res["resumed_from_epoch"] = int(_st.get("current_epoch", _st.get("epoch", 0)))
+            except Exception:
+                pass
             if not args.nnu_skip_prep:
                 print("[B] --resume without --nnu-skip-prep: plan_and_preprocess "
                       "still runs with --clean (the trained checkpoint is kept)",
@@ -499,7 +505,8 @@ def bench_B(args, train_rows, test_rows):
         t0 = time.time()
         N.run(bcmd, os.path.join(d("logs/nnu"), "bench_train.log"), env=env, retries=0)
         res["train_s"] = round(time.time() - t0, 1)
-        res["train_s_per_epoch"] = round(res["train_s"] / max(args.train_epochs, 1), 1)
+        n_trained = max(args.train_epochs - int(res.get("resumed_from_epoch", 0)), 1)
+        res["train_s_per_epoch"] = round(res["train_s"] / n_trained, 1)
         res["train_vram_mb"] = round(poller.max_mib * 1.048576, 1) if poller.max_mib else None
         poller.stop()
         res["train_vram_note"] = "nvidia-smi GPU memory.used peak during nnUNetv2_train sub-process"