logging rework temp save
This commit is contained in:
rpotter6298
@@ -0,0 +1,461 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
V2-parity metadata-only runner.
|
||||
|
||||
Goal:
|
||||
- Match V2HyperTower single-model metadata-only behavior as closely as possible.
|
||||
- Avoid image tower/image IO overhead in forward/training.
|
||||
|
||||
How parity is achieved:
|
||||
- Uses PatientFirstSplitManager (same split policy).
|
||||
- Uses PAPILA profile builders + V2 filters:
|
||||
- eye_train = filter_eye_samples(...)
|
||||
- bilat_val/test = filter_bilateral_samples(...)
|
||||
- Uses V2 training/eval helpers directly:
|
||||
- train_single_epoch(...)
|
||||
- collect_probs_single_components(...)
|
||||
- Uses bridge_mode="metadata_only".
|
||||
|
||||
Implementation detail:
|
||||
- Batch dictionaries still include image slots to satisfy shared V2 helpers,
|
||||
but these are tiny dummy tensors and are never consumed in metadata-only mode.
|
||||
|
||||
Outputs:
|
||||
analysis_data/{run_name}/{eval_mode}/{tower_mode}/fold{N}/
|
||||
y_true.npy
|
||||
probs_classic.npy or probs_ensemble.npy
|
||||
y_true_holdout.npy
|
||||
probs_classic_holdout.npy or probs_ensemble_holdout.npy
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
# ensure repo root is on sys.path when run directly
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(_REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(_REPO_ROOT))
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
|
||||
from classes.v2.bridges import Bridge
|
||||
from classes.v2.loader_factory import (
|
||||
build_balanced_sampler,
|
||||
filter_bilateral_samples,
|
||||
filter_eye_samples,
|
||||
)
|
||||
from classes.v2.metrics import _score_arrays
|
||||
from classes.v2.models import collect_probs_single_components, train_single_epoch
|
||||
from classes.v2.papila_builders import build_papila_data
|
||||
from classes.v2.profiles import build_papila_profile
|
||||
from classes.v2.split_manager import PatientFirstSplitManager
|
||||
from classes.v2.towers import MDTower
|
||||
from classes.v2.utils import choose_device, seed_everything
|
||||
|
||||
|
||||
class MetadataOnlySingleHT(nn.Module):
|
||||
"""SingleEyeHT-compatible shell without real image tower usage."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
clinical_data,
|
||||
num_classes: int,
|
||||
md_hidden_dim: int,
|
||||
fusion_dim: int,
|
||||
dropout: float,
|
||||
use_se: bool,
|
||||
se_reduction: int,
|
||||
se_pre_norm: bool,
|
||||
):
|
||||
super().__init__()
|
||||
# Placeholder module to satisfy phase toggling logic.
|
||||
self.img_tower = nn.Identity()
|
||||
self.md_tower = MDTower(
|
||||
clinical_data=clinical_data,
|
||||
hidden_dim=md_hidden_dim,
|
||||
dropout=dropout,
|
||||
use_se=use_se,
|
||||
se_reduction=se_reduction,
|
||||
se_pre_norm=se_pre_norm,
|
||||
)
|
||||
# img_dim is irrelevant in metadata_only mode, but Bridge defines img head params.
|
||||
self.bridge = Bridge(
|
||||
img_dim=1,
|
||||
meta_dim=self.md_tower.out_dim,
|
||||
num_classes=num_classes,
|
||||
fusion_dim=fusion_dim,
|
||||
mode="metadata_only",
|
||||
use_se=False,
|
||||
se_reduction=16,
|
||||
se_pre_norm=True,
|
||||
)
|
||||
|
||||
|
||||
class EyeMetaDataset(Dataset):
|
||||
"""Eye-level dataset for V2 train_single_epoch input contract."""
|
||||
|
||||
def __init__(self, samples: list[dict]):
|
||||
self.samples = samples
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
|
||||
s = self.samples[idx]
|
||||
return {
|
||||
"image_1": torch.zeros(1, dtype=torch.float32),
|
||||
"matrix_1": torch.as_tensor(s["matrix_1"], dtype=torch.float32),
|
||||
"label_1": torch.tensor(int(s["label_1"]), dtype=torch.long),
|
||||
}
|
||||
|
||||
|
||||
class BilatMetaDataset(Dataset):
|
||||
"""Patient-level bilateral dataset for collect_probs_single_components."""
|
||||
|
||||
def __init__(self, samples: list[dict]):
|
||||
self.samples = samples
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
|
||||
s = self.samples[idx]
|
||||
return {
|
||||
"image_1": torch.zeros(1, dtype=torch.float32),
|
||||
"image_2": torch.zeros(1, dtype=torch.float32),
|
||||
"matrix_1": torch.as_tensor(s["matrix_1"], dtype=torch.float32),
|
||||
"matrix_2": torch.as_tensor(s["matrix_2"], dtype=torch.float32),
|
||||
"label_1": torch.tensor(int(s["label_1"]), dtype=torch.long),
|
||||
}
|
||||
|
||||
|
||||
def _phase_for_epoch(epoch_idx: int, warm_tower: int, warm_fused: int, main_epochs: int) -> tuple[str, int]:
|
||||
if epoch_idx < warm_tower:
|
||||
return "tower_warmup", 0
|
||||
if epoch_idx < (warm_tower + warm_fused):
|
||||
return "fused_warmup", 0
|
||||
if epoch_idx < (warm_tower + warm_fused + main_epochs):
|
||||
main_ep = epoch_idx - warm_tower - warm_fused + 1
|
||||
return "main", main_ep
|
||||
return "done", main_epochs
|
||||
|
||||
|
||||
def _evaluate_single(
|
||||
model: nn.Module,
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
num_classes: int,
|
||||
aggregate_patient: bool,
|
||||
) -> tuple[np.ndarray, np.ndarray, float, float]:
|
||||
y, p_fused, _, p_md = collect_probs_single_components(
|
||||
model, loader, device, aggregate_patient=aggregate_patient
|
||||
)
|
||||
# In metadata_only mode p_fused == p_md; keep md explicitly for clarity.
|
||||
probs = p_md if p_md.size else p_fused
|
||||
acc, auc, _ = _score_arrays(y, probs, num_classes)
|
||||
return y, probs, float(auc), float(acc)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser(description="V2-parity metadata-only runner.")
|
||||
ap.add_argument("--eval-mode", required=True, choices=["binary", "multiclass"])
|
||||
ap.add_argument("--tower-mode", default="single", choices=["single", "ensemble"])
|
||||
ap.add_argument("--run-name", required=True)
|
||||
|
||||
ap.add_argument("--epochs", type=int, default=40, help="Main-phase epochs.")
|
||||
ap.add_argument("--warmup-tower-epochs", type=int, default=None)
|
||||
ap.add_argument("--warmup-fused-epochs", type=int, default=None)
|
||||
ap.add_argument("--batch-size", type=int, default=8)
|
||||
ap.add_argument("--lr", type=float, default=1e-4)
|
||||
ap.add_argument("--weight-decay", type=float, default=0.0)
|
||||
ap.add_argument("--bcd-prob", type=float, default=0.5)
|
||||
|
||||
ap.add_argument("--md-hidden-dim", type=int, default=128)
|
||||
ap.add_argument("--fusion-dim", type=int, default=256)
|
||||
ap.add_argument("--dropout", type=float, default=0.1)
|
||||
ap.add_argument("--use-se", action="store_true")
|
||||
ap.add_argument("--se-reduction", type=int, default=16)
|
||||
ap.add_argument("--se-pre-norm", action="store_true")
|
||||
|
||||
ap.add_argument("--n-splits", type=int, default=5)
|
||||
ap.add_argument("--holdout-per-class", type=int, default=5)
|
||||
ap.add_argument("--holdout-seed", type=int, default=123)
|
||||
ap.add_argument("--fold-seed", type=int, default=42)
|
||||
ap.add_argument("--seed", type=int, default=1234)
|
||||
ap.add_argument("--balanced-sampling", action=argparse.BooleanOptionalAction, default=False)
|
||||
|
||||
ap.add_argument("--analysis-dir", default="analysis_data")
|
||||
ap.add_argument("--image-dir", default="Papila/FundusImages")
|
||||
ap.add_argument("--clinical-dir", default="Papila/ClinicalData")
|
||||
ap.add_argument("--label-col", default="Diagnosis")
|
||||
ap.add_argument("--patient-col", default="Patient ID")
|
||||
ap.add_argument("--cat-cols", nargs="*", default=["Gender", "Phakic/Pseudophakic"])
|
||||
|
||||
args = ap.parse_args()
|
||||
|
||||
seed_everything(args.seed)
|
||||
device = choose_device(None)
|
||||
print(f"Device: {device}", flush=True)
|
||||
|
||||
print("Loading PAPILA data...", flush=True)
|
||||
data = build_papila_data(
|
||||
image_dir=args.image_dir,
|
||||
clinical_dir=args.clinical_dir,
|
||||
label_col=args.label_col,
|
||||
cat_cols=list(args.cat_cols),
|
||||
n_splits=args.n_splits,
|
||||
random_seed=args.fold_seed,
|
||||
)
|
||||
print(f"Loaded: {len(data.df)} rows feature_dim={data.feature_dim}", flush=True)
|
||||
|
||||
num_classes = 2 if args.eval_mode == "binary" else 3
|
||||
df_mode = data.df.copy()
|
||||
if args.eval_mode == "binary":
|
||||
df_mode = df_mode[df_mode[args.label_col].isin([0, 1])].reset_index(drop=True)
|
||||
print(f"[{args.eval_mode}] rows={len(df_mode)}", flush=True)
|
||||
|
||||
class _ClinicalShim:
|
||||
label_col = args.label_col
|
||||
|
||||
def __init__(self, df):
|
||||
self.df = df
|
||||
|
||||
split_args = SimpleNamespace(
|
||||
eval_mode=args.eval_mode,
|
||||
holdout_per_class=args.holdout_per_class,
|
||||
holdout_seed=args.holdout_seed,
|
||||
n_splits=args.n_splits,
|
||||
fold_seed=args.fold_seed,
|
||||
)
|
||||
splitter = PatientFirstSplitManager(patient_col=args.patient_col, label_col=args.label_col)
|
||||
plans = splitter.build_plans(clinical=_ClinicalShim(df_mode), args=split_args)
|
||||
|
||||
profile_eye = build_papila_profile(
|
||||
patient_col=args.patient_col,
|
||||
label_col=args.label_col,
|
||||
sample_mode="eye",
|
||||
)
|
||||
profile_patient = build_papila_profile(
|
||||
patient_col=args.patient_col,
|
||||
label_col=args.label_col,
|
||||
sample_mode="patient",
|
||||
)
|
||||
|
||||
warm_tower = int(args.warmup_tower_epochs) if args.warmup_tower_epochs is not None else 2
|
||||
warm_fused = int(args.warmup_fused_epochs) if args.warmup_fused_epochs is not None else 2
|
||||
total_epochs = warm_tower + warm_fused + int(args.epochs)
|
||||
|
||||
out_root = Path(args.analysis_dir) / args.run_name / args.eval_mode / args.tower_mode
|
||||
out_root.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
fold_metrics = []
|
||||
aggregate_patient = args.tower_mode == "ensemble"
|
||||
|
||||
for fold_idx, split in enumerate(plans[: args.n_splits]):
|
||||
fold_dir = out_root / f"fold{fold_idx}"
|
||||
fold_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
eye_train = filter_eye_samples(profile_eye.build_samples(df=split.train, clinical=data))
|
||||
bilat_val = filter_bilateral_samples(profile_patient.build_samples(df=split.val, clinical=data))
|
||||
|
||||
holdout_bilat = []
|
||||
if split.holdout is not None and not split.holdout.empty:
|
||||
holdout_bilat = filter_bilateral_samples(
|
||||
profile_patient.build_samples(df=split.holdout, clinical=data)
|
||||
)
|
||||
|
||||
if not eye_train or not bilat_val:
|
||||
print(f"[fold {fold_idx+1}] skipped (eye_train={len(eye_train)} bilat_val={len(bilat_val)})", flush=True)
|
||||
fold_metrics.append(
|
||||
{
|
||||
"fold": fold_idx,
|
||||
"best_epoch": None,
|
||||
"best_phase": None,
|
||||
"val_auc": float("nan"),
|
||||
"val_acc": float("nan"),
|
||||
"hld_auc": float("nan"),
|
||||
"hld_acc": float("nan"),
|
||||
"eye_train_n": len(eye_train),
|
||||
"bilat_val_n": len(bilat_val),
|
||||
"bilat_holdout_n": len(holdout_bilat),
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
sampler = build_balanced_sampler(eye_train) if args.balanced_sampling else None
|
||||
train_loader = DataLoader(
|
||||
EyeMetaDataset(eye_train),
|
||||
batch_size=args.batch_size,
|
||||
shuffle=(sampler is None),
|
||||
sampler=sampler,
|
||||
)
|
||||
val_loader = DataLoader(BilatMetaDataset(bilat_val), batch_size=args.batch_size, shuffle=False)
|
||||
holdout_loader = (
|
||||
DataLoader(BilatMetaDataset(holdout_bilat), batch_size=args.batch_size, shuffle=False)
|
||||
if holdout_bilat
|
||||
else None
|
||||
)
|
||||
|
||||
model = MetadataOnlySingleHT(
|
||||
clinical_data=data,
|
||||
num_classes=num_classes,
|
||||
md_hidden_dim=args.md_hidden_dim,
|
||||
fusion_dim=args.fusion_dim,
|
||||
dropout=args.dropout,
|
||||
use_se=bool(args.use_se),
|
||||
se_reduction=int(args.se_reduction),
|
||||
se_pre_norm=bool(args.se_pre_norm),
|
||||
).to(device)
|
||||
opt = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
|
||||
|
||||
best_auc = -1.0
|
||||
best_epoch = 0
|
||||
best_phase = ""
|
||||
best_state = None
|
||||
|
||||
print(
|
||||
f"\n[fold {fold_idx+1}/{args.n_splits}] "
|
||||
f"eye_train_n={len(eye_train)} bilat_val_n={len(bilat_val)} "
|
||||
f"holdout_n={len(holdout_bilat)} warmup={warm_tower}+{warm_fused} total={total_epochs}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
for ep in range(total_epochs):
|
||||
phase, main_ep = _phase_for_epoch(ep, warm_tower, warm_fused, int(args.epochs))
|
||||
tr_loss, tr_acc = train_single_epoch(
|
||||
model,
|
||||
train_loader,
|
||||
opt,
|
||||
device,
|
||||
phase=phase,
|
||||
bcd_prob=float(args.bcd_prob),
|
||||
)
|
||||
|
||||
_, p_val, val_auc, val_acc = _evaluate_single(
|
||||
model,
|
||||
val_loader,
|
||||
device,
|
||||
num_classes,
|
||||
aggregate_patient=aggregate_patient,
|
||||
)
|
||||
|
||||
is_main = phase == "main"
|
||||
if is_main and (not np.isnan(val_auc)) and val_auc > best_auc:
|
||||
best_auc = float(val_auc)
|
||||
best_state = copy.deepcopy(model.state_dict())
|
||||
best_epoch = ep + 1
|
||||
best_phase = phase
|
||||
|
||||
if ep == 0 or (ep + 1) % 10 == 0 or (ep + 1) == total_epochs:
|
||||
print(
|
||||
f" ep {ep+1:>3}/{total_epochs} [{phase}:{main_ep}/{args.epochs}] "
|
||||
f"loss={tr_loss:.4f} acc={tr_acc:.4f} "
|
||||
f"val_auc={val_auc:.4f} val_acc={val_acc:.4f} "
|
||||
f"best_auc={best_auc:.4f}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
if best_state is not None:
|
||||
model.load_state_dict(best_state)
|
||||
|
||||
y_val, p_val, val_auc, val_acc = _evaluate_single(
|
||||
model,
|
||||
val_loader,
|
||||
device,
|
||||
num_classes,
|
||||
aggregate_patient=aggregate_patient,
|
||||
)
|
||||
|
||||
if args.tower_mode == "single":
|
||||
probs_name = "probs_classic.npy"
|
||||
probs_h_name = "probs_classic_holdout.npy"
|
||||
else:
|
||||
probs_name = "probs_ensemble.npy"
|
||||
probs_h_name = "probs_ensemble_holdout.npy"
|
||||
|
||||
np.save(fold_dir / "y_true.npy", y_val)
|
||||
np.save(fold_dir / probs_name, p_val)
|
||||
|
||||
hld_auc = float("nan")
|
||||
hld_acc = float("nan")
|
||||
if holdout_loader is not None:
|
||||
y_h, p_h, hld_auc, hld_acc = _evaluate_single(
|
||||
model,
|
||||
holdout_loader,
|
||||
device,
|
||||
num_classes,
|
||||
aggregate_patient=aggregate_patient,
|
||||
)
|
||||
np.save(fold_dir / "y_true_holdout.npy", y_h)
|
||||
np.save(fold_dir / probs_h_name, p_h)
|
||||
|
||||
print(
|
||||
f" [fold {fold_idx+1}] best_epoch={best_epoch} best_auc={best_auc:.4f} "
|
||||
f"val_auc={val_auc:.4f} val_acc={val_acc:.4f} "
|
||||
f"hld_auc={hld_auc:.4f} hld_acc={hld_acc:.4f}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
fold_metrics.append(
|
||||
{
|
||||
"fold": fold_idx,
|
||||
"best_epoch": best_epoch,
|
||||
"best_phase": best_phase,
|
||||
"val_auc": float(val_auc),
|
||||
"val_acc": float(val_acc),
|
||||
"hld_auc": float(hld_auc),
|
||||
"hld_acc": float(hld_acc),
|
||||
"eye_train_n": len(eye_train),
|
||||
"bilat_val_n": len(bilat_val),
|
||||
"bilat_holdout_n": len(holdout_bilat),
|
||||
}
|
||||
)
|
||||
|
||||
val_aucs = [m["val_auc"] for m in fold_metrics if not np.isnan(m["val_auc"])]
|
||||
hld_aucs = [m["hld_auc"] for m in fold_metrics if not np.isnan(m["hld_auc"])]
|
||||
if val_aucs:
|
||||
print(f"\nMean val AUC: {np.mean(val_aucs):.4f} ± {np.std(val_aucs):.4f}", flush=True)
|
||||
if hld_aucs:
|
||||
print(f"Mean hld AUC: {np.mean(hld_aucs):.4f} ± {np.std(hld_aucs):.4f}", flush=True)
|
||||
|
||||
summary = {
|
||||
"run_name": args.run_name,
|
||||
"eval_mode": args.eval_mode,
|
||||
"tower_mode": args.tower_mode,
|
||||
"bridge_mode": "metadata_only",
|
||||
"model": "MetadataOnlySingleHT",
|
||||
"epochs": int(args.epochs),
|
||||
"warmup_tower_epochs": warm_tower,
|
||||
"warmup_fused_epochs": warm_fused,
|
||||
"md_hidden_dim": int(args.md_hidden_dim),
|
||||
"fusion_dim": int(args.fusion_dim),
|
||||
"dropout": float(args.dropout),
|
||||
"lr": float(args.lr),
|
||||
"weight_decay": float(args.weight_decay),
|
||||
"bcd_prob": float(args.bcd_prob),
|
||||
"balanced_sampling": bool(args.balanced_sampling),
|
||||
"feature_dim": int(data.feature_dim),
|
||||
"timestamp": time.strftime("%Y%m%d_%H%M%S"),
|
||||
"fold_metrics": fold_metrics,
|
||||
"val_auc_mean": float(np.mean(val_aucs)) if val_aucs else None,
|
||||
"val_auc_std": float(np.std(val_aucs)) if val_aucs else None,
|
||||
"hld_auc_mean": float(np.mean(hld_aucs)) if hld_aucs else None,
|
||||
"hld_auc_std": float(np.std(hld_aucs)) if hld_aucs else None,
|
||||
}
|
||||
(out_root / "summary.json").write_text(json.dumps(summary, indent=2))
|
||||
print(f"\nOutputs written to: {out_root}", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user