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.
2026-07-03 08:51:44 +02:00
parent 3d954a4606
commit 708fbc70ce
52 changed files with 2223 additions and 218 deletions
@@ -0,0 +1,107 @@
+"""Memory probe: EfficientNetV2-M at 480x480, bilateral forward+backward, AMP bf16.
+
+Goal: confirm bs=8 fits in 16 GB on the available GPU before committing to the
+full sensitivity experiment.
+
+Mimics the bilateral training step (image tower run twice on OD + OS with shared
+weights, plus a small downstream head + CE loss + Adam step). The clinical tower
+and L1 bridge are omitted; their memory footprint is negligible against V2-M
+activations. Synthetic inputs of the correct shape — no v4 dataset needed.
+
+Run: python v4/scripts/experiments/sensitivity/probe_v2m_480_amp.py
+Expected output: GPU name, peak memory at each phase, fit/OOM verdict.
+"""
+from __future__ import annotations
+
+import sys
+import torch
+from torch import nn
+from torchvision.models import efficientnet_v2_m
+
+
+BATCH = 8
+RES = 480
+DTYPE = torch.bfloat16
+
+
+def fmt_gb(bytes_):
+    return f"{bytes_ / 1024**3:.2f} GB"
+
+
+def main():
+    if not torch.cuda.is_available():
+        print("No CUDA/ROCm device available; probe requires a GPU.", file=sys.stderr)
+        sys.exit(2)
+
+    device = torch.device("cuda")
+    gpu_name = torch.cuda.get_device_name(0)
+    gpu_total = torch.cuda.get_device_properties(0).total_memory
+    print(f"GPU: {gpu_name}   total VRAM: {fmt_gb(gpu_total)}")
+    print(f"Config: V2-M, bs={BATCH}, res={RES}, bilateral 2x forward, AMP={DTYPE}")
+    print("-" * 70)
+
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+
+    backbone = efficientnet_v2_m(weights=None).to(device)
+    feat_dim = backbone.classifier[1].in_features
+    backbone.classifier = nn.Identity()
+    head = nn.Sequential(
+        nn.LayerNorm(feat_dim),
+        nn.Linear(feat_dim, 256),
+        nn.GELU(),
+        nn.Linear(256, 2),
+    ).to(device)
+    opt = torch.optim.Adam(
+        list(backbone.parameters()) + list(head.parameters()),
+        lr=1e-4,
+    )
+
+    print(f"After model + optimizer load: allocated={fmt_gb(torch.cuda.memory_allocated())}  "
+          f"peak={fmt_gb(torch.cuda.max_memory_allocated())}")
+
+    x_od = torch.randn(BATCH, 3, RES, RES, device=device)
+    x_os = torch.randn(BATCH, 3, RES, RES, device=device)
+    y = torch.randint(0, 2, (BATCH,), device=device)
+
+    print(f"After synthetic inputs:        allocated={fmt_gb(torch.cuda.memory_allocated())}  "
+          f"peak={fmt_gb(torch.cuda.max_memory_allocated())}")
+
+    try:
+        opt.zero_grad(set_to_none=True)
+        with torch.autocast(device_type="cuda", dtype=DTYPE):
+            z_od = backbone(x_od)
+            z_os = backbone(x_os)
+            z = z_od + z_os
+            logits = head(z)
+            loss = nn.functional.cross_entropy(logits, y)
+
+        print(f"After forward:                 allocated={fmt_gb(torch.cuda.memory_allocated())}  "
+              f"peak={fmt_gb(torch.cuda.max_memory_allocated())}")
+
+        loss.backward()
+        print(f"After backward:                allocated={fmt_gb(torch.cuda.memory_allocated())}  "
+              f"peak={fmt_gb(torch.cuda.max_memory_allocated())}")
+
+        opt.step()
+        print(f"After optimizer step:          allocated={fmt_gb(torch.cuda.memory_allocated())}  "
+              f"peak={fmt_gb(torch.cuda.max_memory_allocated())}")
+
+        torch.cuda.synchronize()
+        peak = torch.cuda.max_memory_allocated()
+        headroom = gpu_total - peak
+        print("-" * 70)
+        print(f"VERDICT: FIT  | peak={fmt_gb(peak)} of {fmt_gb(gpu_total)}  "
+              f"headroom={fmt_gb(headroom)} ({100 * headroom / gpu_total:.1f}%)")
+        print(f"Loss value: {loss.item():.4f}")
+
+    except torch.cuda.OutOfMemoryError as e:
+        peak = torch.cuda.max_memory_allocated()
+        print("-" * 70)
+        print(f"VERDICT: OOM  | peak before OOM={fmt_gb(peak)} of {fmt_gb(gpu_total)}")
+        print(f"OOM details: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()