Add distributed server implementation and protocol definitions

- Introduced `protocol.py` for shared data models used in server/client communication, including request and response schemas for registration, job submission, and status updates.
- Implemented `server.py` to manage a SQLite job queue and client registry, handling job polling, status updates, and job completion.
- Created a cheat sheet for server usage, detailing commands for starting the server, submitting jobs, and monitoring clients.
- Added several experiment configuration files for various training setups, including geometry vector injections and baseline ensembles.

v4/classes/stages/fusion.py

@@ -11,7 +11,8 @@ import torch.nn.functional as F
 from v4.classes.dataset import LoaderShell, to_label_tensor
 from v4.classes.metrics import score_arrays, compute_extended_metrics, tune_binary_threshold
 from v4.classes.stages.helpers import (
-    encode_embedding, get_out_dim, resolve_input_dims, phase_for_epoch,
+    class_weights_from_shell, encode_embedding, get_out_dim, resolve_input_dims,
+    phase_for_epoch,
 )
 
 
@@ -25,15 +26,15 @@ def collect_probs(
     loader,
     device,
     num_classes:  int,
-) -> tuple[np.ndarray, np.ndarray]:
-    """Eval pass for one fusion stage; returns (y_true, softmax_probs)."""
+) -> tuple[np.ndarray, np.ndarray, list, np.ndarray]:
+    """Eval pass for one fusion stage; returns (y_true, softmax_probs, entity_ids, embeddings)."""
     from v4.classes.dataset import to_label_tensor
     bridge.eval(); primary_head.eval()
     for t in towers.values():
         t.eval()
     inputs       = stage_cfg["inputs"]
     is_bilateral = isinstance(inputs, dict)
-    y_all, p_all = [], []
+    y_all, p_all, ids_all, z_all = [], [], [], []
 
     with torch.no_grad():
         for batch in loader:
@@ -57,10 +58,14 @@ def collect_probs(
             logits = primary_head(z)
             y_all.append(to_label_tensor(y, device).cpu().numpy())
             p_all.append(F.softmax(logits, dim=1).cpu().numpy())
+            z_all.append(z.cpu().numpy())
+            ids_all.extend(batch.get("entity_id", []))
 
     if not y_all:
-        return np.zeros(0, dtype=np.int64), np.zeros((0, num_classes), dtype=np.float32)
-    return np.concatenate(y_all), np.concatenate(p_all, axis=0)
+        return (np.zeros(0, dtype=np.int64), np.zeros((0, num_classes), dtype=np.float32),
+                [], np.zeros((0, 0), dtype=np.float32))
+    return (np.concatenate(y_all), np.concatenate(p_all, axis=0),
+            ids_all, np.concatenate(z_all, axis=0))
 
 
 def run(
@@ -88,9 +93,13 @@ def run(
     inputs       = stage_cfg["inputs"]
     is_bilateral = isinstance(inputs, dict)
 
-    s_train = data.build_shells(split.train, level=level, label_filter=label_filter)
-    s_val   = data.build_shells(split.val,   level=level, label_filter=label_filter)
-    s_test  = (data.build_shells(split.test, level=level, label_filter=label_filter)
+    eye_filter = stage_cfg.get("eye_filter", None)
+    s_train = data.build_shells(split.train, level=level, label_filter=label_filter,
+                                eye_filter=eye_filter)
+    s_val   = data.build_shells(split.val,   level=level, label_filter=label_filter,
+                                eye_filter=eye_filter)
+    s_test  = (data.build_shells(split.test, level=level, label_filter=label_filter,
+                                 eye_filter=eye_filter)
                if split.test is not None else LoaderShell(entries=[]))
 
     if not s_val.entries:
@@ -116,7 +125,11 @@ def run(
         h_dim = get_out_dim(hs["input"], towers, {**stage_models, name: bridge})
         h_mod = importlib.import_module(hs.get("module", "v4.classes.heads.classifier"))
         h_cls = getattr(h_mod, hs.get("class", "ClassificationHead"))
-        head_models[hs["name"]] = h_cls(h_dim, num_classes).to(device)
+        existing = stage_models.get(hs["name"])
+        head_models[hs["name"]] = (
+            existing.to(device) if existing is not None
+            else h_cls(h_dim, num_classes, **hs.get("args", {})).to(device)
+        )
 
     primary_hs_cfg = next((hs for hs in head_stage_cfgs if not hs.get("bcd", False)), None)
     bcd_head_cfgs  = [hs for hs in head_stage_cfgs if hs.get("bcd", False)]
@@ -132,11 +145,23 @@ def run(
         for p in m.parameters():
             p.requires_grad_(False)
         m.eval()
+    for h in head_models.values():
+        for p in h.parameters():
+            p.requires_grad_(True)
 
     # ── Optimizer ────────────────────────────────────────────────────────────
     train_towers = stage_cfg.get("train_towers", False)
+    if train_towers:
+        # Only train towers that are direct inputs to this stage (not all towers globally).
+        # For nt_od with inputs ["img_od", "cd_od"] this trains only those two; other
+        # eye's towers remain untouched.
+        direct_inputs = list(inputs.values()) if isinstance(inputs, dict) else inputs
+        tower_params  = [p for n in direct_inputs if n in towers
+                         for p in towers[n].parameters()]
+    else:
+        tower_params = []
     opt_params = (
-        ([p for t in towers.values() for p in t.parameters()] if train_towers else []) +
+        tower_params +
         list(bridge.parameters()) +
         [p for h in head_models.values() for p in h.parameters()]
     )
@@ -147,6 +172,17 @@ def run(
     wf         = 0 if is_bilateral else warmup_cfg.get("fused_epochs",  0)
     bcd_prob   = cfg["training"].get("bcd_prob", 0.5)
 
+    cw = class_weights_from_shell(
+        s_train, num_classes, device,
+        enabled=cfg["training"].get("class_weighted", False),
+    )
+    if cw is not None:
+        print(
+            f"  fold{fold+1} [{name}] class weights: "
+            + ", ".join(f"{i}={w:.3f}" for i, w in enumerate(cw.tolist())),
+            flush=True,
+        )
+
     # ── Epoch loop ────────────────────────────────────────────────────────────
     for epoch in range(epochs):
         bridge.train()
@@ -200,7 +236,7 @@ def run(
             if is_bilateral or phase == "fused_warmup":
                 logits = head_logits.get(primary_hs_cfg["name"])
             elif phase == "tower_warmup" and bcd_head_cfgs:
-                losses = [F.cross_entropy(head_logits[hs["name"]], y_t)
+                losses = [F.cross_entropy(head_logits[hs["name"]], y_t, weight=cw)
                           for hs in bcd_head_cfgs if hs["name"] in head_logits]
                 if not losses:
                     continue
@@ -217,7 +253,7 @@ def run(
 
             if logits is None:
                 continue
-            loss = F.cross_entropy(logits, y_t)
+            loss = F.cross_entropy(logits, y_t, weight=cw)
             opt.zero_grad(); loss.backward(); opt.step()
             total_correct += int((logits.argmax(1) == y_t).sum())
             total_loss    += loss.item() * len(y_t)
@@ -226,8 +262,8 @@ def run(
         tr_loss = total_loss / total_n if total_n else nan
         tr_acc  = total_correct / total_n if total_n else nan
 
-        y_v, p_v = collect_probs(bridge, primary_head, stage_cfg, towers,
-                                  stage_models, cfg_stages, val_loader, device, num_classes)
+        y_v, p_v, _, _ = collect_probs(bridge, primary_head, stage_cfg, towers,
+                                       stage_models, cfg_stages, val_loader, device, num_classes)
         _, val_auc, _ = score_arrays(y_v, p_v, num_classes) if y_v.size else (nan, nan, nan)
         print(
             f"  fold{fold+1} [{name}] ep{epoch+1:03d}/{epochs} [{phase:14s}]"
@@ -236,8 +272,8 @@ def run(
         )
 
     # ── Final eval ────────────────────────────────────────────────────────────
-    y_val, p_val = collect_probs(bridge, primary_head, stage_cfg, towers,
-                                  stage_models, cfg_stages, val_loader, device, num_classes)
+    y_val, p_val, ids_val, z_val = collect_probs(bridge, primary_head, stage_cfg, towers,
+                                                  stage_models, cfg_stages, val_loader, device, num_classes)
     val_acc, val_auc, val_n = (score_arrays(y_val, p_val, num_classes)
                                if y_val.size else (nan, nan, nan))
     ext = compute_extended_metrics(y_val, p_val, num_classes) if y_val.size else {}
@@ -247,10 +283,11 @@ def run(
             and num_classes == 2 and y_val.size >= 2):
         val_threshold = tune_binary_threshold(y_val, p_val[:, 1])
 
+    y_te = p_te = ids_te = z_te = None
     test_auc = test_acc = test_n = nan
     if test_loader is not None:
-        y_te, p_te = collect_probs(bridge, primary_head, stage_cfg, towers,
-                                    stage_models, cfg_stages, test_loader, device, num_classes)
+        y_te, p_te, ids_te, z_te = collect_probs(bridge, primary_head, stage_cfg, towers,
+                                                  stage_models, cfg_stages, test_loader, device, num_classes)
         test_acc, test_auc, test_n = (score_arrays(y_te, p_te, num_classes)
                                       if y_te.size else (nan, nan, nan))
 
@@ -270,4 +307,10 @@ def run(
         f"{name}_test_acc":      test_acc,
         f"{name}_test_n":        test_n,
     }
-    return updated, metrics
+    pred_data = {
+        name: {
+            "val_y": y_val, "val_p": p_val, "val_ids": ids_val, "val_z": z_val,
+            "test_y": y_te, "test_p": p_te, "test_ids": ids_te, "test_z": z_te,
+        }
+    }
+    return updated, metrics, pred_data

v4/classes/stages/helpers.py

@@ -1,9 +1,32 @@
 """stages/helpers — shared utilities for stage runners."""
 from __future__ import annotations
 
+from collections import Counter
+
 import torch
 
 
+def class_weights_from_shell(
+    shell, num_classes: int, device, *, enabled: bool = True
+) -> torch.Tensor | None:
+    """Return inverse-frequency CE weights normalised to mean 1, or None.
+
+    weights[i] = (N_total / num_classes) / N_class_i  → rare classes weighted higher.
+    Mean(weights) ≈ 1 so overall loss magnitude is unchanged.
+
+    Classes absent from the shell get weight 1.0 (no division-by-zero).
+    """
+    if not enabled or shell is None or not shell.entries:
+        return None
+    counts  = Counter(int(e.label) for e in shell.entries)
+    n_total = sum(counts.values())
+    weights = []
+    for c in range(num_classes):
+        n_c = counts.get(c, 0)
+        weights.append(1.0 if n_c == 0 else n_total / (num_classes * n_c))
+    return torch.tensor(weights, dtype=torch.float32, device=device)
+
+
 def get_out_dim(name: str, towers: dict, stage_models: dict) -> int:
     if name in towers:
         return towers[name].out_dim

v4/classes/stages/parallel.py

@@ -0,0 +1,394 @@
+"""stages/parallel — runs multiple same-type sub-stages in a shared epoch loop.
+
+Used to train bilateral pairs (OD + OS) simultaneously rather than sequentially.
+Sub-stages must all be the same type: either all 'warm' or all 'fusion'.
+"""
+from __future__ import annotations
+
+import importlib
+from random import choice, random as _random
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from v4.classes.dataset import LoaderShell, to_label_tensor
+from v4.classes.metrics import score_arrays, compute_extended_metrics, tune_binary_threshold
+from v4.classes.stages.helpers import (
+    class_weights_from_shell, get_out_dim, resolve_input_dims, phase_for_epoch,
+)
+from v4.classes.stages.fusion import collect_probs
+
+
+def run(
+    stage_cfg:        dict,
+    cfg:              dict,
+    towers:           dict,
+    stage_models:     dict,
+    data,
+    split,
+    label_filter,
+    num_classes:      int,
+    device,
+    fold:             int,
+    cfg_stages:       list[dict],
+    _make_loader,
+    _balanced_sampler,
+) -> tuple[dict, dict, dict]:
+    sub_cfgs  = stage_cfg["stages"]
+    sub_types = {s["type"] for s in sub_cfgs}
+
+    if sub_types == {"warm"}:
+        return _parallel_warm(sub_cfgs, cfg, towers, stage_models, data, split,
+                              label_filter, num_classes, device, fold,
+                              cfg_stages, _make_loader, _balanced_sampler)
+    elif sub_types == {"fusion"}:
+        return _parallel_fusion(sub_cfgs, cfg, towers, stage_models, data, split,
+                                label_filter, num_classes, device, fold,
+                                cfg_stages, _make_loader)
+    else:
+        raise ValueError(
+            f"parallel stage sub-stages must all be the same type (warm or fusion), "
+            f"got {sub_types}"
+        )
+
+
+# ── Parallel warm ─────────────────────────────────────────────────────────────
+
+def _parallel_warm(
+    sub_cfgs, cfg, towers, stage_models, data, split,
+    label_filter, num_classes, device, fold, cfg_stages, _make_loader, _balanced_sampler,
+):
+    bs = cfg["training"]["batch_size"]
+
+    contexts = []
+    for sc in sub_cfgs:
+        tower_name   = sc["tower"]
+        level        = sc["level"]
+        shell_filter = sc.get("shell_filter", {})
+        n_epochs     = sc.get("epochs", 0)
+
+        s_train = data.build_shells(split.train, level=level, label_filter=label_filter,
+                                    **shell_filter)
+        loader  = _make_loader(s_train, {tower_name: towers[tower_name]},
+                               batch_size=bs, shuffle=False,
+                               sampler=_balanced_sampler(s_train))
+        head_name = sc.get("head_name")
+        if head_name:
+            head_cfg = next((s for s in cfg_stages if s.get("name") == head_name), None)
+            if head_cfg is None:
+                raise ValueError(f"warm stage requested head_name={head_name!r}, but no such head exists")
+            h_mod = importlib.import_module(head_cfg.get("module", "v4.classes.heads.classifier"))
+            h_cls = getattr(h_mod, head_cfg.get("class", "ClassificationHead"))
+            probe = stage_models.get(head_name)
+            if probe is None:
+                probe = h_cls(towers[tower_name].out_dim, num_classes, **head_cfg.get("args", {}))
+            probe = probe.to(device)
+        else:
+            probe = torch.nn.Linear(towers[tower_name].out_dim, num_classes).to(device)
+        opt   = torch.optim.Adam(
+            list(towers[tower_name].parameters()) + list(probe.parameters()),
+            lr=cfg["training"]["lr"],
+        )
+        cw = class_weights_from_shell(
+            s_train, num_classes, device,
+            enabled=cfg["training"].get("class_weighted", False),
+        )
+        contexts.append({
+            "name": tower_name, "n_epochs": n_epochs,
+            "loader": loader, "probe": probe, "opt": opt,
+            "head_name": head_name, "class_weights": cw,
+        })
+
+    active = {c["name"] for c in contexts if c["n_epochs"] > 0}
+    for name, t in towers.items():
+        for p in t.parameters():
+            p.requires_grad_(name in active)
+
+    max_epochs = max((c["n_epochs"] for c in contexts), default=0)
+    for epoch in range(max_epochs):
+        for ctx in contexts:
+            if epoch >= ctx["n_epochs"]:
+                continue
+            tower_name = ctx["name"]
+            towers[tower_name].train()
+            total_loss = total_correct = total_n = 0
+            for batch in ctx["loader"]:
+                y = batch.get("label")
+                x = batch.get(tower_name)
+                if not torch.is_tensor(y) or not torch.is_tensor(x):
+                    continue
+                y_t    = to_label_tensor(y, device)
+                logits = ctx["probe"](towers[tower_name](x.to(device)))
+                loss   = F.cross_entropy(logits, y_t, weight=ctx["class_weights"])
+                ctx["opt"].zero_grad(); loss.backward(); ctx["opt"].step()
+                total_loss    += loss.item() * len(y_t)
+                total_correct += int((logits.argmax(1) == y_t).sum())
+                total_n       += len(y_t)
+            if total_n:
+                print(
+                    f"  fold{fold+1} [warm/{tower_name}]"
+                    f" ep{epoch+1:03d}/{ctx['n_epochs']}"
+                    f"  loss={total_loss/total_n:.4f}  acc={total_correct/total_n:.3f}",
+                    flush=True,
+                )
+
+    for t in towers.values():
+        for p in t.parameters():
+            p.requires_grad_(True)
+
+    updated = dict(stage_models)
+    for ctx in contexts:
+        if ctx.get("head_name"):
+            updated[ctx["head_name"]] = ctx["probe"]
+
+    return updated, {}, {}
+
+
+# ── Parallel fusion ───────────────────────────────────────────────────────────
+
+def _parallel_fusion(
+    sub_cfgs, cfg, towers, stage_models, data, split,
+    label_filter, num_classes, device, fold, cfg_stages, _make_loader,
+):
+    nan      = float("nan")
+    bs       = cfg["training"]["batch_size"]
+    bcd_prob = cfg["training"].get("bcd_prob", 0.5)
+
+    # Freeze all prior stage models once, before building any bridges.
+    for m in stage_models.values():
+        for p in m.parameters():
+            p.requires_grad_(False)
+        m.eval()
+
+    # ── Per-sub-stage setup ───────────────────────────────────────────────────
+    contexts = []
+    for sc in sub_cfgs:
+        name       = sc["name"]
+        level      = sc["level"]
+        inputs     = sc["inputs"]
+        epochs     = sc["epochs"]
+        eye_filter = sc.get("eye_filter", None)
+
+        s_train = data.build_shells(split.train, level=level, label_filter=label_filter,
+                                    eye_filter=eye_filter)
+        s_val   = data.build_shells(split.val,   level=level, label_filter=label_filter,
+                                    eye_filter=eye_filter)
+        s_test  = (data.build_shells(split.test, level=level, label_filter=label_filter,
+                                     eye_filter=eye_filter)
+                   if split.test is not None else LoaderShell(entries=[]))
+
+        if not s_val.entries:
+            print(f"  fold{fold+1}: no val samples for stage {name!r}, skipping.", flush=True)
+            continue
+
+        input_dims = resolve_input_dims(inputs, towers, stage_models)
+        bmod       = importlib.import_module(sc["module"])
+        bridge     = getattr(bmod, sc["class"])(input_dims, **sc.get("args", {})).to(device)
+
+        # Head stages for this sub-stage.
+        head_stage_cfgs = [s for s in cfg_stages if s["type"] == "head"
+                           and s.get("train_with") == name]
+        head_models: dict[str, torch.nn.Module] = {}
+        for hs in head_stage_cfgs:
+            h_dim = get_out_dim(hs["input"], towers, {**stage_models, name: bridge})
+            h_mod = importlib.import_module(hs.get("module", "v4.classes.heads.classifier"))
+            h_cls = getattr(h_mod, hs.get("class", "ClassificationHead"))
+            existing = stage_models.get(hs["name"])
+            head_models[hs["name"]] = (
+                existing.to(device) if existing is not None
+                else h_cls(h_dim, num_classes, **hs.get("args", {})).to(device)
+            )
+        for h in head_models.values():
+            for p in h.parameters():
+                p.requires_grad_(True)
+
+        primary_hs_cfg = next((hs for hs in head_stage_cfgs if not hs.get("bcd", False)), None)
+        bcd_head_cfgs  = [hs for hs in head_stage_cfgs if hs.get("bcd", False)]
+
+        if primary_hs_cfg is None:
+            print(f"  WARNING: no primary head for stage {name!r}; skipping.", flush=True)
+            continue
+
+        primary_head = head_models[primary_hs_cfg["name"]]
+
+        train_towers  = sc.get("train_towers", False)
+        direct_inputs = inputs if isinstance(inputs, list) else list(inputs.values())
+        tower_params  = ([p for n in direct_inputs if n in towers
+                          for p in towers[n].parameters()]
+                         if train_towers else [])
+        opt_params    = (tower_params + list(bridge.parameters()) +
+                         [p for h in head_models.values() for p in h.parameters()])
+        opt = torch.optim.Adam(opt_params, lr=cfg["training"]["lr"])
+
+        warmup_cfg = sc.get("warmup", {})
+        wt         = warmup_cfg.get("tower_epochs", 0)
+        wf         = warmup_cfg.get("fused_epochs", 0)
+
+        train_loader = _make_loader(s_train, towers, batch_size=bs, shuffle=True)
+        val_loader   = _make_loader(s_val,   towers, batch_size=bs, shuffle=False)
+        test_loader  = (_make_loader(s_test, towers, batch_size=bs, shuffle=False)
+                        if s_test.entries else None)
+
+        cw = class_weights_from_shell(
+            s_train, num_classes, device,
+            enabled=cfg["training"].get("class_weighted", False),
+        )
+        contexts.append({
+            "name": name, "inputs": inputs, "epochs": epochs,
+            "bridge": bridge, "head_models": head_models,
+            "primary_head": primary_head, "primary_hs_cfg": primary_hs_cfg,
+            "bcd_head_cfgs": bcd_head_cfgs,
+            "opt": opt, "wt": wt, "wf": wf,
+            "train_towers": train_towers, "direct_inputs": direct_inputs,
+            "train_loader": train_loader, "val_loader": val_loader,
+            "test_loader": test_loader, "sc": sc,
+            "class_weights": cw,
+        })
+
+    if not contexts:
+        return stage_models, {}, {}
+
+    max_epochs = max(c["epochs"] for c in contexts)
+
+    # ── Shared epoch loop ─────────────────────────────────────────────────────
+    for epoch in range(max_epochs):
+        for ctx in contexts:
+            if epoch >= ctx["epochs"]:
+                continue
+
+            name   = ctx["name"]
+            bridge = ctx["bridge"]
+            inputs = ctx["inputs"]
+            wt, wf = ctx["wt"], ctx["wf"]
+            phase  = phase_for_epoch(epoch, wt, wf)
+
+            bridge.train()
+            for h in ctx["head_models"].values():
+                h.train()
+            if ctx["train_towers"]:
+                for n in ctx["direct_inputs"]:
+                    if n in towers:
+                        towers[n].train()
+            else:
+                for n in ctx["direct_inputs"]:
+                    if n in towers:
+                        towers[n].eval()
+
+            if hasattr(bridge, "set_phase"):
+                bridge.set_phase(phase)
+
+            total_loss = total_correct = total_n = 0
+
+            for batch in ctx["train_loader"]:
+                y = batch.get("label")
+                if not torch.is_tensor(y):
+                    continue
+                y_t = to_label_tensor(y, device)
+                if y_t.numel() == 0:
+                    continue
+
+                local_embs = {n: towers[n](batch[n].to(device)) for n in inputs
+                              if n in batch and torch.is_tensor(batch[n])}
+                if len(local_embs) != len(inputs):
+                    continue
+                local_embs[name] = bridge(list(local_embs[n] for n in inputs))
+
+                head_logits = {
+                    hs["name"]: ctx["head_models"][hs["name"]](local_embs[hs["input"]])
+                    for hs in ([ctx["primary_hs_cfg"]] + ctx["bcd_head_cfgs"])
+                    if hs["input"] in local_embs
+                }
+
+                if phase == "fused_warmup":
+                    logits = head_logits.get(ctx["primary_hs_cfg"]["name"])
+                elif phase == "tower_warmup" and ctx["bcd_head_cfgs"]:
+                    losses = [F.cross_entropy(head_logits[hs["name"]], y_t,
+                                              weight=ctx["class_weights"])
+                              for hs in ctx["bcd_head_cfgs"] if hs["name"] in head_logits]
+                    if not losses:
+                        continue
+                    loss = sum(losses) / len(losses)
+                    ctx["opt"].zero_grad(); loss.backward(); ctx["opt"].step()
+                    total_loss += loss.item() * len(y_t)
+                    total_n    += len(y_t)
+                    continue
+                else:
+                    if ctx["bcd_head_cfgs"] and _random() < bcd_prob:
+                        logits = head_logits.get(choice(ctx["bcd_head_cfgs"])["name"])
+                    else:
+                        logits = head_logits.get(ctx["primary_hs_cfg"]["name"])
+
+                if logits is None:
+                    continue
+                loss = F.cross_entropy(logits, y_t, weight=ctx["class_weights"])
+                ctx["opt"].zero_grad(); loss.backward(); ctx["opt"].step()
+                total_correct += int((logits.argmax(1) == y_t).sum())
+                total_loss    += loss.item() * len(y_t)
+                total_n       += len(y_t)
+
+            tr_loss = total_loss / total_n if total_n else nan
+            tr_acc  = total_correct / total_n if total_n else nan
+
+            y_v, p_v, _, _ = collect_probs(bridge, ctx["primary_head"], ctx["sc"], towers,
+                                           stage_models, cfg_stages, ctx["val_loader"],
+                                           device, num_classes)
+            _, val_auc, _ = score_arrays(y_v, p_v, num_classes) if y_v.size else (nan, nan, nan)
+            print(
+                f"  fold{fold+1} [{name}] ep{epoch+1:03d}/{ctx['epochs']} [{phase:14s}]"
+                f"  loss={tr_loss:.4f}  acc={tr_acc:.3f}  val_auc={val_auc:.4f}",
+                flush=True,
+            )
+
+    # ── Final eval + collect results ──────────────────────────────────────────
+    updated    = dict(stage_models)
+    all_metrics: dict = {}
+    all_preds:   dict = {}
+
+    for ctx in contexts:
+        name         = ctx["name"]
+        bridge       = ctx["bridge"]
+        primary_head = ctx["primary_head"]
+
+        updated[name] = bridge
+        updated.update(ctx["head_models"])
+
+        y_val, p_val, ids_val, z_val = collect_probs(bridge, primary_head, ctx["sc"], towers,
+                                                      stage_models, cfg_stages, ctx["val_loader"],
+                                                      device, num_classes)
+        val_acc, val_auc, val_n = (score_arrays(y_val, p_val, num_classes)
+                                   if y_val.size else (nan, nan, nan))
+        ext = compute_extended_metrics(y_val, p_val, num_classes) if y_val.size else {}
+
+        val_threshold = 0.5
+        if (cfg["training"].get("tune_binary_threshold")
+                and num_classes == 2 and y_val.size >= 2):
+            val_threshold = tune_binary_threshold(y_val, p_val[:, 1])
+
+        y_te = p_te = ids_te = z_te = None
+        test_auc = test_acc = test_n = nan
+        if ctx["test_loader"] is not None:
+            y_te, p_te, ids_te, z_te = collect_probs(bridge, primary_head, ctx["sc"], towers,
+                                                      stage_models, cfg_stages, ctx["test_loader"],
+                                                      device, num_classes)
+            test_acc, test_auc, test_n = (score_arrays(y_te, p_te, num_classes)
+                                          if y_te.size else (nan, nan, nan))
+
+        all_metrics.update({
+            f"{name}_val_auc":       val_auc,
+            f"{name}_val_acc":       val_acc,
+            f"{name}_val_n":         val_n,
+            f"{name}_val_kappa":     ext.get("kappa", nan),
+            f"{name}_val_mcc":       ext.get("mcc", nan),
+            f"{name}_val_f1":        ext.get("macro_f1", nan),
+            f"{name}_val_threshold": val_threshold,
+            f"{name}_test_auc":      test_auc,
+            f"{name}_test_acc":      test_acc,
+            f"{name}_test_n":        test_n,
+        })
+        all_preds[name] = {
+            "val_y": y_val, "val_p": p_val, "val_ids": ids_val, "val_z": z_val,
+            "test_y": y_te, "test_p": p_te, "test_ids": ids_te, "test_z": z_te,
+        }
+
+    return updated, all_metrics, all_preds

v4/classes/stages/warm.py

@@ -1,11 +1,13 @@
-"""stages/warm — warm stage runner: pre-trains a single tower with a temporary probe."""
+"""stages/warm — warm stage runner: pre-trains a single tower and optional real head."""
 from __future__ import annotations
 
+import importlib
+
 import torch
 import torch.nn.functional as F
 
 from v4.classes.dataset import to_label_tensor
-from v4.classes.stages.helpers import phase_for_epoch
+from v4.classes.stages.helpers import class_weights_from_shell
 
 
 def run(
@@ -20,16 +22,27 @@ def run(
     fold:        int,
     _make_loader,
     _balanced_sampler,
-) -> None:
-    """Pre-train one tower using a temporary linear probe (probe discarded after)."""
-    tower_name = stage_cfg["tower"]
-    n_epochs   = stage_cfg.get("epochs", 0)
-    level      = stage_cfg["level"]
+    stage_models: dict | None = None,
+    cfg_stages:   list[dict] | None = None,
+) -> dict:
+    """Pre-train one tower.
+
+    If ``head_name`` is set on the stage config, train that real downstream head
+    and return it in ``stage_models``. Otherwise, fall back to a temporary linear
+    probe for backward-compatible representation warmup.
+    """
+    tower_name   = stage_cfg["tower"]
+    n_epochs     = stage_cfg.get("epochs", 0)
+    level        = stage_cfg["level"]
+    shell_filter = stage_cfg.get("shell_filter", {})
+    stage_models = dict(stage_models or {})
+    cfg_stages   = list(cfg_stages or [])
 
     if n_epochs == 0:
-        return
+        return stage_models
 
-    s_train = data.build_shells(split.train, level=level, label_filter=label_filter)
+    s_train = data.build_shells(split.train, level=level, label_filter=label_filter,
+                                **shell_filter)
     bs      = cfg["training"]["batch_size"]
     loader  = _make_loader(
         s_train, {tower_name: towers[tower_name]},
@@ -41,13 +54,32 @@ def run(
         for p in t.parameters():
             p.requires_grad_(n == tower_name)
 
-    probe = torch.nn.Linear(towers[tower_name].out_dim, num_classes).to(device)
+    head_name = stage_cfg.get("head_name")
+    if head_name:
+        head_cfg = next((s for s in cfg_stages if s.get("name") == head_name), None)
+        if head_cfg is None:
+            raise ValueError(f"warm stage requested head_name={head_name!r}, but no such head exists")
+        h_mod = importlib.import_module(head_cfg.get("module", "v4.classes.heads.classifier"))
+        h_cls = getattr(h_mod, head_cfg.get("class", "ClassificationHead"))
+        probe = stage_models.get(head_name)
+        if probe is None:
+            probe = h_cls(towers[tower_name].out_dim, num_classes, **head_cfg.get("args", {}))
+        probe = probe.to(device)
+    else:
+        probe = torch.nn.Linear(towers[tower_name].out_dim, num_classes).to(device)
+
     opt   = torch.optim.Adam(
         list(towers[tower_name].parameters()) + list(probe.parameters()),
         lr=cfg["training"]["lr"],
     )
 
+    cw = class_weights_from_shell(
+        s_train, num_classes, device,
+        enabled=cfg["training"].get("class_weighted", False),
+    )
+
     towers[tower_name].train()
+    probe.train()
     for epoch in range(n_epochs):
         total_loss = total_correct = total_n = 0
         for batch in loader:
@@ -57,7 +89,7 @@ def run(
                 continue
             y_t    = to_label_tensor(y, device)
             logits = probe(towers[tower_name](x.to(device)))
-            loss   = F.cross_entropy(logits, y_t)
+            loss   = F.cross_entropy(logits, y_t, weight=cw)
             opt.zero_grad(); loss.backward(); opt.step()
             total_loss    += loss.item() * len(y_t)
             total_correct += int((logits.argmax(1) == y_t).sum())
@@ -71,3 +103,7 @@ def run(
     for t in towers.values():
         for p in t.parameters():
             p.requires_grad_(True)
+
+    if head_name:
+        stage_models[head_name] = probe
+    return stage_models