ryan/hypertower
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added fused classifier head

Browse Source
This commit is contained in:
rpotter6298
2026-03-02 06:44:22 +01:00
parent fd1452c187
commit 080b5999fa
11 changed files with 583 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
classes/v2/loader_factory.py
+13 -2
View File
@@ -4,7 +4,7 @@ from dataclasses import dataclass
from typing import Any, Callable, Optional
import torch
from torch.utils.data import DataLoader
from torch.utils.data import DataLoader, WeightedRandomSampler
from .network_manager import LoaderBundle, PatientSplit
from .slot_dataset import SlotDataset, slot_collate
@@ -201,6 +201,7 @@ def make_loader(
batch_size: int,
shuffle: bool,
num_workers: int,
sampler: Optional[WeightedRandomSampler] = None,
) -> DataLoader:
ds = SlotDataset(
samples,
@@ -211,12 +212,22 @@ def make_loader(
return DataLoader(
ds,
batch_size=batch_size,
shuffle=shuffle,
shuffle=(shuffle if sampler is None else False),
sampler=sampler,
num_workers=num_workers,
collate_fn=slot_collate,
)
def build_balanced_sampler(samples: list[dict], label_key: str = "label_1") -> WeightedRandomSampler:
"""Return a WeightedRandomSampler that equalises class frequency for training."""
from collections import Counter
labels = [s[label_key] for s in samples]
counts = Counter(labels)
weights = [1.0 / counts[lbl] for lbl in labels]
return WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)
def to_label_tensor(labels, device: torch.device) -> torch.Tensor:
if torch.is_tensor(labels):
return labels.to(device=device, dtype=torch.long)
classes/v2/models.py
+106
View File
@@ -157,6 +157,53 @@ class BilateralHT(nn.Module):
return out_f
class FusedEnsembleHT(nn.Module):
"""
SingleEyeHT base with a per-eye attention scorer for bilateral fusion.
The base model is trained eye-level (identical to ensemble mode).
After base training completes, the base is frozen and only the
eye_scorer is trained on bilateral (patient-level) samples.
At inference, eye_scorer is applied independently to each eye's logit
vector to produce a scalar attention score. Softmax over the two scores
gives attention weights; the final logit is a weighted sum:
score_od = eye_scorer(logit_od) # [B, 1]
score_os = eye_scorer(logit_os) # [B, 1]
alpha = softmax([score_od, score_os]) # [B, 2], sums to 1
out = alpha[:,0:1]*logit_od + alpha[:,1:2]*logit_os
Because eye_scorer is applied to each eye with the same weights, the
mechanism is permutation-equivariant — there is no left/right positional
bias. Through training on bilateral labels the scorer learns to give high
scores to logits that point strongly toward the GC class, creating the
desired asymmetry: a confidently GC eye dominates the patient prediction
more than a comparably confident healthy eye would.
"""
def __init__(self, base: SingleEyeHT, num_classes: int):
super().__init__()
self.base = base
# Applied independently to each eye's logit → scalar attention score.
# Learns the GC-direction in logit space from bilateral labels.
self.eye_scorer = nn.Linear(num_classes, 1, bias=True)
def forward(
self,
x_od: torch.Tensor,
meta_od: torch.Tensor,
x_os: torch.Tensor,
meta_os: torch.Tensor,
) -> torch.Tensor:
logit_od = self.base(x_od, meta_od) # [B, C]
logit_os = self.base(x_os, meta_os) # [B, C]
scores = torch.cat([self.eye_scorer(logit_od),
self.eye_scorer(logit_os)], dim=1) # [B, 2]
alpha = torch.softmax(scores, dim=1) # [B, 2]
return alpha[:, 0:1] * logit_od + alpha[:, 1:2] * logit_os # [B, C]
# ---------------------------------------------------------------------------
# Phase control
# ---------------------------------------------------------------------------
@@ -327,6 +374,39 @@ def train_bilateral_epoch(
)
def train_fusion_epoch(
model: FusedEnsembleHT,
loader: DataLoader,
opt,
device: torch.device,
) -> tuple[float, float]:
"""Train only the fusion head; the base SingleEyeHT is frozen in eval mode."""
model.base.eval()
model.eye_scorer.train()
total_loss = total_correct = total_n = 0
for batch in loader:
x1 = batch.get("image_1"); m1 = batch.get("matrix_1")
x2 = batch.get("image_2"); m2 = batch.get("matrix_2")
y = batch.get("label_1")
if not (torch.is_tensor(x1) and torch.is_tensor(m1) and
torch.is_tensor(x2) and torch.is_tensor(m2)):
continue
y_t = _to_label_tensor(y, device)
out = model(x1.to(device), m1.to(device), x2.to(device), m2.to(device))
loss = F.cross_entropy(out, y_t)
opt.zero_grad()
loss.backward()
opt.step()
bs = y_t.shape[0]
total_loss += float(loss.item()) * bs
total_correct += int((out.argmax(1) == y_t).sum())
total_n += bs
return (
total_loss / total_n if total_n else float("nan"),
total_correct / total_n if total_n else float("nan"),
)
# ---------------------------------------------------------------------------
# Inference helpers
# ---------------------------------------------------------------------------
@@ -420,6 +500,32 @@ def collect_probs_bilateral(
return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0)
def collect_probs_fused(
model: FusedEnsembleHT,
loader: DataLoader,
device: torch.device,
) -> tuple[np.ndarray, np.ndarray]:
"""Patient-level fused-head eval. Returns (y_true [N], probs [N, C])."""
model.eval()
y_chunks, p_chunks = [], []
with torch.no_grad():
for batch in loader:
x1 = batch.get("image_1"); m1 = batch.get("matrix_1")
x2 = batch.get("image_2"); m2 = batch.get("matrix_2")
y = batch.get("label_1")
if not (torch.is_tensor(x1) and torch.is_tensor(m1) and
torch.is_tensor(x2) and torch.is_tensor(m2)):
continue
y_t = _to_label_tensor(y, device)
p = F.softmax(model(x1.to(device), m1.to(device),
x2.to(device), m2.to(device)), dim=1)
y_chunks.append(y_t.cpu().numpy())
p_chunks.append(p.cpu().numpy())
if not y_chunks:
return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32)
return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0)
def collect_probs_single_components(
model: SingleEyeHT,
loader: DataLoader,
classes/v2/results.py
+16 -1
View File
@@ -1,7 +1,7 @@
"""Result dataclasses and serialisation helpers for V2 fold outputs."""
from __future__ import annotations
from dataclasses import dataclass
from dataclasses import dataclass, field
from typing import Optional
import numpy as np
@@ -84,6 +84,19 @@ class FoldResult:
# Training sample counts
single_train_n: int
bilat_train_n: int
# Fused head (ensemble + --fused-head; nan / None if --fused-head not used)
fused_val_auc: float = float("nan")
fused_val_acc: float = float("nan")
fused_val_kappa: float = float("nan")
fused_val_mcc: float = float("nan")
fused_val_f1: float = float("nan")
fused_val_recall: Optional[str] = None
fused_val_ece: float = float("nan")
fused_val_threshold: float = float("nan")
fused_val_bias: Optional[str] = None
fused_val_n: int = 0
fused_holdout_auc: float = float("nan")
fused_holdout_acc: float = float("nan")
@dataclass
@@ -94,3 +107,5 @@ class FoldArtifacts:
probs_ensemble: Optional[np.ndarray]
y_true_bilat: Optional[np.ndarray]
probs_bilat: Optional[np.ndarray]
y_true_fused: Optional[np.ndarray] = None
probs_fused: Optional[np.ndarray] = None
classes/v2/v2_hypertower.py
+142 -5
View File
@@ -24,6 +24,7 @@ import torch
from classes.v2.croppers import build_image_preprocessor_from_args
from classes.v2.dataset import _ClinicalView # noqa: F401 (re-exported for compat)
from classes.v2.loader_factory import (
build_balanced_sampler,
filter_bilateral_samples,
filter_eye_samples,
make_loader,
@@ -31,14 +32,17 @@ from classes.v2.loader_factory import (
from classes.v2.metrics import _score_arrays, _svf, _tune_and_snap
from classes.v2.models import (
BilateralHT,
FusedEnsembleHT,
SingleEyeHT,
V2ModeComparisonOps,
collect_probs_bilateral,
collect_probs_bilateral_components,
collect_probs_classic,
collect_probs_ensemble,
collect_probs_fused,
collect_probs_single_components,
train_bilateral_epoch,
train_fusion_epoch,
train_single_epoch,
)
from classes.v2.papila_builders import build_papila_data
@@ -118,6 +122,8 @@ class V2HyperTower:
ap.add_argument("--backbone", default="refugelike")
ap.add_argument("--freeze-ratio", type=float, default=0.0)
ap.add_argument("--augment", action="store_true")
ap.add_argument("--balanced-sampling", action="store_true",
help="Use WeightedRandomSampler during training to equalise class frequency (default: off).")
ap.add_argument("--num-workers", type=int, default=0)
ap.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
ap.add_argument("--seed", type=int, default=1234)
@@ -186,6 +192,16 @@ class V2HyperTower:
ap.add_argument("--log-every", type=int, default=1)
ap.add_argument("--save-checkpoints", action=argparse.BooleanOptionalAction, default=True,
help="Save best_single.pt / best_holdout_single.pt per fold (use --no-save-checkpoints to disable)")
ap.add_argument(
"--fused-head", action="store_true",
help="(ensemble mode only) After base SingleEyeHT training, freeze it and train a "
"small logit-level MLP fusion head on bilateral samples instead of averaging "
"OD/OS softmax probabilities.",
)
ap.add_argument(
"--fusion-epochs", type=int, default=10,
help="Number of epochs to train the fusion head (--fused-head, ensemble mode only).",
)
return ap
# ------------------------------------------------------------------
@@ -414,7 +430,8 @@ class V2HyperTower:
nan = _nan()
tower_mode = "single" if tower_mode == "classic" else tower_mode
run_single = tower_mode in ("single", "ensemble")
run_bilat = tower_mode == "bilateral"
run_bilat = tower_mode == "bilateral"
run_fused = (tower_mode == "ensemble") and bool(getattr(args, "fused_head", False))
global_warmup_tower = getattr(args, "warmup_tower_epochs", None)
global_warmup_fused = getattr(args, "warmup_fused_epochs", None)
@@ -506,22 +523,38 @@ class V2HyperTower:
loader_kw = dict(batch_size=args.batch_size, num_workers=args.num_workers)
# ---- loaders ---------------------------------------------------
use_balanced = bool(getattr(args, "balanced_sampling", False))
train_single_loader = None
train_bilat_loader = None
if run_single:
single_sampler = build_balanced_sampler(eye_train) if use_balanced else None
train_single_loader = make_loader(
eye_train, slots_eye,
image_transform=single.transform,
image_preprocessor=image_preprocessor,
shuffle=True,
sampler=single_sampler,
**loader_kw,
)
if run_bilat:
bilat_sampler = build_balanced_sampler(bilat_train) if use_balanced else None
train_bilat_loader = make_loader(
bilat_train, slots_patient,
image_transform=bilateral.transform,
image_preprocessor=image_preprocessor,
shuffle=True,
sampler=bilat_sampler,
**loader_kw,
)
elif run_fused:
# Fused head trains on bilateral samples using the single model's transform.
fused_sampler = build_balanced_sampler(bilat_train) if use_balanced else None
train_bilat_loader = make_loader(
bilat_train, slots_patient,
image_transform=single.transform,
image_preprocessor=image_preprocessor,
shuffle=True,
sampler=fused_sampler,
**loader_kw,
)
eval_transform = build_eval_transform(args.backbone)
@@ -579,11 +612,13 @@ class V2HyperTower:
best_epoch_bilat = 0
best_single_state: Optional[dict] = None
best_bilat_state: Optional[dict] = None
snap_classic: dict = {}
snap_ensemble: dict = {}
snap_bilat: dict = {}
snap_classic: dict = {}
snap_ensemble: dict = {}
snap_bilat: dict = {}
snap_holdout_single: dict = {}
snap_holdout_bilat: dict = {}
snap_fused: dict = {}
snap_holdout_fused: dict = {}
best_holdout_single_auc = -1.0
best_holdout_bilat_auc = -1.0
best_epoch_holdout_single = 0
@@ -863,6 +898,73 @@ class V2HyperTower:
if best_holdout_bilat_state is not None:
torch.save(best_holdout_bilat_state, fold_dir / "best_holdout_bilateral.pt")
# ---- Phase 2: fused head training (ensemble + --fused-head only) ----
best_fused_auc = -1.0
best_fused_state: Optional[dict] = None
best_holdout_fused_auc = -1.0
if run_fused and best_single_state is not None:
# Revert base to its best val checkpoint, then freeze it.
single.load_state_dict(best_single_state)
for p in single.parameters():
p.requires_grad_(False)
fused = FusedEnsembleHT(single, num_classes).to(device)
opt_fused = torch.optim.Adam(fused.eye_scorer.parameters(), lr=args.lr)
fusion_epochs = int(getattr(args, "fusion_epochs", 10))
print(
f" [fold {fold+1}] Phase 2: training fusion head "
f"bilat_train_n={len(bilat_train)} fusion_epochs={fusion_epochs}",
flush=True,
)
for fep in range(fusion_epochs):
fu_loss, fu_acc = train_fusion_epoch(fused, train_bilat_loader, opt_fused, device)
y_fu, p_fu = collect_probs_fused(fused, val_loader, device)
fu_auc = _score_arrays(y_fu, p_fu, num_classes)[1]
# Holdout eval (if available)
fu_hld_auc = nan
if holdout_loader is not None:
y_fu_h, p_fu_h = collect_probs_fused(fused, holdout_loader, device)
fu_hld_auc = _score_arrays(y_fu_h, p_fu_h, num_classes)[1]
is_best_fused = not np.isnan(fu_auc) and fu_auc > best_fused_auc
if is_best_fused:
best_fused_auc = fu_auc
best_fused_state = copy.deepcopy(fused.state_dict())
fu_acc_val = _score_arrays(y_fu, p_fu, num_classes)[0]
snap_fused, _, _, _ = _tune_and_snap(y_fu, p_fu, fu_acc_val, num_classes, args, args.ece_bins)
is_best_hld_fused = not np.isnan(fu_hld_auc) and fu_hld_auc > best_holdout_fused_auc
if is_best_hld_fused:
best_holdout_fused_auc = fu_hld_auc
snap_holdout_fused = {"auc": fu_hld_auc, "acc": _score_arrays(y_fu_h, p_fu_h, num_classes)[0]}
if (fep + 1) % max(1, getattr(args, "log_every", 1)) == 0:
print(
f" [fold {fold+1}] fusion ep{fep+1:>3} "
f"loss={fu_loss:.4f} train_acc={fu_acc:.4f} "
f"val_auc={fu_auc:.4f} hld_auc={fu_hld_auc:.4f}"
f"{' *' if is_best_fused else ''}",
flush=True,
)
if best_fused_state is not None:
if args.save_checkpoints:
torch.save(best_fused_state, fold_dir / "best_fused.pt")
print(
f" [fold {fold+1}] BEST "
f"fused_head(acc={snap_fused.get('acc', nan):.4f},"
f"auc={snap_fused.get('auc', nan):.4f}) "
f"kappa={snap_fused.get('kappa', nan):.4f} "
f"F1={snap_fused.get('macro_f1', nan):.4f} "
f"ECE={snap_fused.get('ece', nan):.4f} "
f"holdout_auc={snap_holdout_fused.get('auc', nan):.4f}",
flush=True,
)
if run_single:
if tower_mode == "single":
print(
@@ -911,6 +1013,11 @@ class V2HyperTower:
else:
y_bi_best = p_bi_best = None
y_fu_best = p_fu_best = None
if run_fused and best_fused_state is not None:
fused.load_state_dict(best_fused_state)
y_fu_best, p_fu_best = collect_probs_fused(fused, val_loader, device)
return FoldResult(
mode=mode, fold=fold,
best_epoch_single=best_epoch_single, best_epoch_bilat=best_epoch_bilat,
@@ -953,10 +1060,23 @@ class V2HyperTower:
holdout_n=len(holdout_bilat),
single_train_n=len(eye_train),
bilat_train_n=len(bilat_train),
fused_val_auc=snap_fused.get("auc", nan),
fused_val_acc=snap_fused.get("acc", nan),
fused_val_kappa=snap_fused.get("kappa", nan),
fused_val_mcc=snap_fused.get("mcc", nan),
fused_val_f1=snap_fused.get("macro_f1", nan),
fused_val_recall=_sv(snap_fused.get("per_class_recall")),
fused_val_ece=snap_fused.get("ece", nan),
fused_val_threshold=snap_fused.get("threshold", nan),
fused_val_bias=_svf(snap_fused.get("bias")),
fused_val_n=snap_fused.get("n", 0),
fused_holdout_auc=snap_holdout_fused.get("auc", nan),
fused_holdout_acc=snap_holdout_fused.get("acc", nan),
), FoldArtifacts(
y_true_classic=y_cl_best, probs_classic=p_cl_best,
y_true_ensemble=y_en_best, probs_ensemble=p_en_best,
y_true_bilat=y_bi_best, probs_bilat=p_bi_best,
y_true_fused=y_fu_best, probs_fused=p_fu_best,
)
# ------------------------------------------------------------------
@@ -976,6 +1096,7 @@ class V2HyperTower:
("classic_best_val", "classic_val"),
("ensemble_best_val", "ensemble_val"),
("bilat_best_val", "bilat_val"),
("fused_best_val", "fused_val"),
]:
sub = {}
for m in ["auc", "acc", "kappa", "mcc", "f1", "ece", "threshold"]:
@@ -990,6 +1111,7 @@ class V2HyperTower:
("classic_holdout", "classic_holdout"),
("ensemble_holdout", "ensemble_holdout"),
("bilat_holdout", "bilat_holdout"),
("fused_holdout", "fused_holdout"),
]:
sub = {}
for m in ["auc", "acc"]:
@@ -1003,6 +1125,7 @@ class V2HyperTower:
for delta_label, prefix_a, prefix_b in [
("delta_ensemble_vs_classic", "classic_val", "ensemble_val"),
("delta_bilat_vs_ensemble", "ensemble_val", "bilat_val"),
("delta_fused_vs_ensemble", "ensemble_val", "fused_val"),
]:
delta = {}
for m in ["auc", "f1", "kappa"]:
@@ -1022,7 +1145,8 @@ class V2HyperTower:
out["eval_note"] = (
"classic=eye-level SingleEyeHT; "
"ensemble=patient-level SingleEyeHT (OD+OS averaged); "
"bilateral=patient-level BilateralHT"
"bilateral=patient-level BilateralHT; "
"fused=ensemble base + learned logit-level fusion head"
)
return out
@@ -1034,8 +1158,11 @@ class V2HyperTower:
cv = s["classic_best_val"]
ev = s["ensemble_best_val"]
bv = s["bilat_best_val"]
fv = s["fused_best_val"]
d1 = s["delta_ensemble_vs_classic"]
d2 = s["delta_bilat_vs_ensemble"]
d3 = s["delta_fused_vs_ensemble"]
has_fused = fv["auc_mean"] is not None
print(f"\n=== Summary [{mode}] — best-epoch val ===")
print(f" {'':26s} {'AUC':>8} {'ACC':>8} {'Kappa':>8} {'F1-mac':>8} {'ECE':>8}")
@@ -1043,6 +1170,8 @@ class V2HyperTower:
rows = [("single (eye-lvl eval)", cv)]
elif tower_mode == "ensemble":
rows = [("ensemble (pat-lvl eval)", ev)]
if has_fused:
rows.append(("fused_head(pat-lvl eval)", fv))
elif tower_mode == "bilateral":
rows = [("bilateral (bilat eval)", bv)]
else:
@@ -1051,6 +1180,8 @@ class V2HyperTower:
("ensemble (pat-lvl eval)", ev),
("bilateral (bilat eval)", bv),
]
if has_fused:
rows.append(("fused_head(pat-lvl eval)", fv))
for label, d in rows:
print(
f" {label:26s} "
@@ -1068,6 +1199,12 @@ class V2HyperTower:
f"{f(d2['auc_mean']):>8} {'':>8} "
f"{f(d2['kappa_mean']):>8} {f(d2['f1_mean']):>8}"
)
if has_fused and tower_mode in ("ensemble", None):
print(
f" {'Δ fused−ensemble':26s} "
f"{f(d3['auc_mean']):>8} {'':>8} "
f"{f(d3['kappa_mean']):>8} {f(d3['f1_mean']):>8}"
)
# ---------------------------------------------------------------------------