rpotter6298
708fbc70ce
- Introduced `bridge_attention_ceiling_check.py` for variance decomposition analysis on bridge attention configurations. - Added `bridge_attention_readout.py` to perform per-tower gate and contribution readouts, including AUC sanity checks. - Created multiple JSON configuration files for backbone replication experiments, including anonymous CV variants and basic backbones. - Implemented sensitivity experiments to evaluate the impact of axial length inclusion and EfficientNetV2-M performance at higher resolutions. - Added a memory probe script to assess GPU memory usage during training with EfficientNetV2-M.
301 lines
12 KiB
Python
301 lines
12 KiB
Python
"""Per-tower gate / contribution readout + per-head AUC sanity check.
|
|
|
|
For each gated_<backbone> run (and ortho_<backbone>), loads each fold's
|
|
checkpoints and runs inference on the test set, capturing in a single pass:
|
|
• per-sample per-stream bridge attention (sigmoid gates for
|
|
GatedAdditiveBridge; ||W_i·z_i|| pre-LN projection norms for OrthoBridge)
|
|
• per-fold AUCs of the eye-level aux heads (img_aux, cd_aux, nt_head) —
|
|
sanity check that the image tower really is getting stronger as the
|
|
backbone improves, and that the clinical tower stays consistent
|
|
|
|
Outputs:
|
|
v4/results/experiments/bridge_attention/_summary/gated_gates.csv
|
|
v4/results/experiments/bridge_attention/_summary/ortho_stream_norms.csv
|
|
Both CSVs include img_aux_auc, cd_aux_auc, nt_head_auc columns.
|
|
|
|
Run:
|
|
python -m v4.scripts.analysis.bridge_attention_readout
|
|
python -m v4.scripts.analysis.bridge_attention_readout --bridges gated
|
|
python -m v4.scripts.analysis.bridge_attention_readout --reps 3 # subsample for speed
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
from tqdm import tqdm
|
|
|
|
from v4.classes.split_manager import SplitManager
|
|
from v4.classes.v4_hypertower import _make_loader, build_towers, load_data
|
|
import v4.classes.v4_hypertower as orch
|
|
|
|
|
|
SWEEP_ROOT = Path("v4/results/experiments/bridge_attention")
|
|
OUT_DIR = SWEEP_ROOT / "_summary"
|
|
|
|
BACKBONES = [
|
|
"mobilenet_v2",
|
|
"resnet50",
|
|
"efficientnet_b0",
|
|
"efficientnet_v2_m",
|
|
"refugelike",
|
|
"refuge_efficientnet_v2_m",
|
|
]
|
|
|
|
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
# Shared fold-module loader (trimmed version of F8's helper)
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
|
|
def _load_fold_bridge(ckpt_dir: Path, cfg: dict, data, device):
|
|
"""Load towers + nt bridge + eye-level aux heads from a single fold ckpt dir."""
|
|
from v4.classes.heads.classifier import ClassificationHead
|
|
|
|
towers = build_towers(cfg["towers"], data)
|
|
for name, tower in towers.items():
|
|
tower.load_state_dict(
|
|
torch.load(ckpt_dir / f"tower_{name}.pt", map_location="cpu")
|
|
)
|
|
tower.to(device).eval()
|
|
|
|
stage_by_name = {s["name"]: s for s in cfg["stages"]}
|
|
nt_cfg = stage_by_name["nt"]
|
|
nt_mod = importlib.import_module(nt_cfg["module"])
|
|
nt = getattr(nt_mod, nt_cfg["class"])(
|
|
[towers[n].out_dim for n in nt_cfg["inputs"]],
|
|
**nt_cfg.get("args", {}),
|
|
).to(device)
|
|
nt.load_state_dict(torch.load(ckpt_dir / "stage_nt.pt", map_location="cpu"))
|
|
nt.eval()
|
|
|
|
nc = cfg["num_classes"]
|
|
img_aux = ClassificationHead(towers["img"].out_dim, nc).to(device)
|
|
cd_aux = ClassificationHead(towers["cd"].out_dim, nc).to(device)
|
|
nt_head = ClassificationHead(nt.out_dim, nc).to(device)
|
|
img_aux.load_state_dict(torch.load(ckpt_dir / "stage_img_aux.pt", map_location="cpu"))
|
|
cd_aux.load_state_dict(torch.load(ckpt_dir / "stage_cd_aux.pt", map_location="cpu"))
|
|
nt_head.load_state_dict(torch.load(ckpt_dir / "stage_nt_head.pt", map_location="cpu"))
|
|
img_aux.eval(); cd_aux.eval(); nt_head.eval()
|
|
aux = {"img_aux": img_aux, "cd_aux": cd_aux, "nt_head": nt_head}
|
|
return towers, nt, nt_cfg, aux
|
|
|
|
|
|
def _splits_from_cfg(cfg: dict, data):
|
|
label_filter = cfg.get("label_filter", None)
|
|
df_mode = data.df.copy()
|
|
if label_filter is not None:
|
|
df_mode = df_mode[df_mode[data.label_col].isin(label_filter)].reset_index(drop=True)
|
|
identity_cols = getattr(data, "identity_cols", [])
|
|
identity_level = cfg.get("split_identity_level", 1)
|
|
group_col = (
|
|
identity_cols[identity_level - 1]
|
|
if identity_level and identity_cols
|
|
else None
|
|
)
|
|
return SplitManager(group_col=group_col).build_plans(
|
|
df_mode,
|
|
label_col=data.label_col,
|
|
n_splits=cfg.get("folds", 5),
|
|
seed=cfg.get("fold_seed", 100),
|
|
), label_filter
|
|
|
|
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
# Unified inference: capture bridge attention + per-head logits in one pass
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
|
|
def _install_attention_hooks(nt, bridge_kind: str):
|
|
"""Register the right hook(s) for the bridge kind. Returns (handles, finalize)
|
|
where finalize() reads back the captured per-sample (N, n_streams) array."""
|
|
if bridge_kind == "gated":
|
|
captured: list[torch.Tensor] = []
|
|
def hook(_module, _input, output):
|
|
captured.append(output.detach().cpu())
|
|
handles = [nt.gate.register_forward_hook(hook)]
|
|
def finalize():
|
|
return (
|
|
torch.cat(captured, dim=0).numpy() if captured else np.empty((0, 0))
|
|
)
|
|
return handles, finalize
|
|
|
|
if bridge_kind == "ortho":
|
|
# Pre-LN projection norm ||W_i z_i||_2 per stream. Post-LN norm is
|
|
# √fusion_dim by construction, so we hook the linear projection itself.
|
|
per_stream: list[list[torch.Tensor]] = [[] for _ in range(len(nt.inner.W))]
|
|
handles = []
|
|
for i, w in enumerate(nt.inner.W):
|
|
def make_hook(idx):
|
|
def hook(_module, _input, output):
|
|
per_stream[idx].append(output.detach().norm(dim=-1).cpu())
|
|
return hook
|
|
handles.append(w.register_forward_hook(make_hook(i)))
|
|
def finalize():
|
|
cols = [
|
|
torch.cat(c, dim=0).numpy() if c else np.array([])
|
|
for c in per_stream
|
|
]
|
|
if not all(s.size for s in cols):
|
|
return np.empty((0, 0))
|
|
return np.stack(cols, axis=1)
|
|
return handles, finalize
|
|
|
|
raise ValueError(f"unknown bridge_kind {bridge_kind!r}")
|
|
|
|
|
|
def extract_one_fold(
|
|
bridge_kind: str,
|
|
towers, nt, aux: dict,
|
|
data, split_obj, label_filter, device,
|
|
):
|
|
"""Single inference pass that returns:
|
|
attention : ndarray (N, n_streams) — bridge-specific attention signal
|
|
head_aucs : dict[str, float] — AUC of softmax[:,1] for each aux head
|
|
n_samples : int
|
|
"""
|
|
from sklearn.metrics import roc_auc_score
|
|
import torch.nn.functional as F
|
|
|
|
if split_obj.test is None:
|
|
return np.empty((0, 0)), {}, 0
|
|
|
|
shell = data.build_shells(
|
|
split_obj.test, level="patient", label_filter=label_filter
|
|
)
|
|
loader = _make_loader(shell, towers, batch_size=8, shuffle=False)
|
|
|
|
handles, finalize = _install_attention_hooks(nt, bridge_kind)
|
|
img_logits, cd_logits, nt_logits, ys = [], [], [], []
|
|
|
|
try:
|
|
with torch.no_grad():
|
|
for batch in loader:
|
|
y = batch["label"].detach().cpu().numpy()
|
|
for side in ("a", "b"):
|
|
z_img = towers["img"](batch["img"][side].to(device))
|
|
z_cd = towers["cd"](batch["cd"][side].to(device))
|
|
img_logits.append(aux["img_aux"](z_img).cpu())
|
|
cd_logits.append(aux["cd_aux"](z_cd).cpu())
|
|
z_fused = nt([z_img, z_cd]) # attention hooks capture here
|
|
nt_logits.append(aux["nt_head"](z_fused).cpu())
|
|
ys.append(y)
|
|
attention = finalize()
|
|
finally:
|
|
for h in handles:
|
|
h.remove()
|
|
|
|
if not ys:
|
|
return attention, {}, 0
|
|
|
|
y_all = np.concatenate(ys)
|
|
head_aucs: dict[str, float] = {}
|
|
for name, buf in [
|
|
("img_aux", img_logits),
|
|
("cd_aux", cd_logits),
|
|
("nt_head", nt_logits),
|
|
]:
|
|
logits = torch.cat(buf, dim=0).numpy()
|
|
probs = (np.exp(logits - logits.max(axis=1, keepdims=True))
|
|
/ np.exp(logits - logits.max(axis=1, keepdims=True))
|
|
.sum(axis=1, keepdims=True))
|
|
try:
|
|
head_aucs[name] = float(roc_auc_score(y_all, probs[:, 1]))
|
|
except Exception:
|
|
head_aucs[name] = float("nan")
|
|
|
|
return attention, head_aucs, int(attention.shape[0])
|
|
|
|
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
# Driver
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
|
|
def run_for_bridge(bridge_kind: str, max_reps: int | None, device) -> pd.DataFrame:
|
|
rows: list[dict] = []
|
|
for backbone in BACKBONES:
|
|
run_dir = SWEEP_ROOT / f"{bridge_kind}_{backbone}"
|
|
if not run_dir.exists():
|
|
print(f" [{bridge_kind}/{backbone}] (missing)")
|
|
continue
|
|
rep_dirs = sorted(run_dir.glob("rep*"))
|
|
if max_reps is not None:
|
|
rep_dirs = rep_dirs[:max_reps]
|
|
for rep_dir in tqdm(rep_dirs, desc=f"{bridge_kind}/{backbone}", unit="rep"):
|
|
summary_path = rep_dir / "binary" / "summary.json"
|
|
if not summary_path.exists():
|
|
continue
|
|
cfg = json.loads(summary_path.read_text())["config"]
|
|
orch.cfg_ref = cfg
|
|
data = load_data(cfg)
|
|
splits, label_filter = _splits_from_cfg(cfg, data)
|
|
|
|
try:
|
|
rep_idx = int(rep_dir.name.replace("rep", ""))
|
|
except ValueError:
|
|
continue
|
|
|
|
for fold_idx in range(cfg.get("folds", 5)):
|
|
ckpt_dir = rep_dir / "binary" / "checkpoints" / f"fold{fold_idx}"
|
|
if not ckpt_dir.exists():
|
|
continue
|
|
towers, nt, _, aux = _load_fold_bridge(ckpt_dir, cfg, data, device)
|
|
arr, head_aucs, _ = extract_one_fold(
|
|
bridge_kind, towers, nt, aux, data,
|
|
splits[fold_idx], label_filter, device,
|
|
)
|
|
|
|
if arr.size:
|
|
means = arr.mean(axis=0)
|
|
rows.append({
|
|
"bridge": bridge_kind,
|
|
"backbone": backbone,
|
|
"rep": rep_idx,
|
|
"fold": fold_idx,
|
|
"n_samples": arr.shape[0],
|
|
"img_value": float(means[0]),
|
|
"cd_value": float(means[1]),
|
|
"img_share": float(means[0] / (means.sum() + 1e-8)),
|
|
"img_aux_auc": head_aucs.get("img_aux", float("nan")),
|
|
"cd_aux_auc": head_aucs.get("cd_aux", float("nan")),
|
|
"nt_head_auc": head_aucs.get("nt_head", float("nan")),
|
|
})
|
|
|
|
del towers, nt, aux
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
return pd.DataFrame(rows)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description=__doc__)
|
|
ap.add_argument(
|
|
"--bridges", default="gated,ortho",
|
|
help="Comma-separated bridge kinds to process (default: gated,ortho)",
|
|
)
|
|
ap.add_argument(
|
|
"--reps", type=int, default=None,
|
|
help="Cap reps per backbone for a quick first pass (default: all)",
|
|
)
|
|
args = ap.parse_args()
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"device: {device}")
|
|
|
|
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
for bridge in args.bridges.split(","):
|
|
bridge = bridge.strip()
|
|
if not bridge:
|
|
continue
|
|
df = run_for_bridge(bridge, args.reps, device)
|
|
suffix = "gates" if bridge == "gated" else "stream_norms"
|
|
out = OUT_DIR / f"{bridge}_{suffix}.csv"
|
|
df.to_csv(out, index=False)
|
|
print(f"saved {out} ({len(df)} rows)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|