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.
108 lines
3.7 KiB
Python
108 lines
3.7 KiB
Python
"""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()
|