Add analysis scripts and experiment configurations for bridge attention and sensitivity studies
- 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.
This commit is contained in:
rpotter6298
@@ -0,0 +1,300 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user