Add t-query logits and hidden export script
This commit is contained in:
325
export_tquery_logits_hidden.py
Normal file
325
export_tquery_logits_hidden.py
Normal file
@@ -0,0 +1,325 @@
|
||||
"""Export landmark risk logits and hidden states for t_query ages.
|
||||
|
||||
This script follows evaluate_event_free_survival.py's data loading,
|
||||
landmark construction, checkpoint loading, and readout logic, but only exports:
|
||||
|
||||
* all token/disease risk logits from ``model.calc_risk(hidden)``;
|
||||
* the corresponding landmark hidden state.
|
||||
|
||||
The two large arrays are saved separately as .npy files. Row metadata is saved
|
||||
as a CSV with matching row order.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from evaluate_auc_v2 import (
|
||||
LandmarkDataset,
|
||||
build_model_from_dataset,
|
||||
cfg_get,
|
||||
load_checkpoint_state_dict,
|
||||
load_json_config,
|
||||
load_model_state,
|
||||
resolve_dist_mode_for_checkpoint,
|
||||
resolve_eval_device,
|
||||
validate_dataset_metadata,
|
||||
)
|
||||
from evaluate_event_free_survival import (
|
||||
IndexedLandmarkDataset,
|
||||
build_first_occurrence_maps_for_landmarks,
|
||||
collate_indexed_landmark_fn,
|
||||
infer_landmark_hidden,
|
||||
load_eval_sequence_dataset,
|
||||
make_landmark_ages,
|
||||
)
|
||||
|
||||
|
||||
def numpy_float_dtype(name: str) -> np.dtype:
|
||||
key = str(name).lower()
|
||||
if key in {"float16", "fp16", "half"}:
|
||||
return np.dtype(np.float16)
|
||||
if key in {"float32", "fp32", "single"}:
|
||||
return np.dtype(np.float32)
|
||||
raise ValueError(f"dtype must be float16 or float32, got {name!r}")
|
||||
|
||||
|
||||
def output_paths_for_run(
|
||||
run_path: Path,
|
||||
eval_split: str,
|
||||
landmark_start: float,
|
||||
landmark_stop: float,
|
||||
landmark_step: float,
|
||||
output_dir: Optional[str],
|
||||
) -> tuple[Path, Path, Path, Path]:
|
||||
suffix = f"{eval_split}_t{landmark_start:g}-{landmark_stop:g}_step{landmark_step:g}"
|
||||
base_dir = Path(output_dir) if output_dir else run_path
|
||||
return (
|
||||
base_dir / f"tquery_logits_{suffix}.npy",
|
||||
base_dir / f"tquery_hidden_{suffix}.npy",
|
||||
base_dir / f"tquery_metadata_{suffix}.csv",
|
||||
base_dir / f"tquery_export_config_{suffix}.json",
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Export landmark risk logits and hidden states for t_query ages."
|
||||
)
|
||||
parser.add_argument("--run_path", type=str, required=True)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory for output files. Defaults to run_path.",
|
||||
)
|
||||
parser.add_argument("--logits_path", type=str, default=None)
|
||||
parser.add_argument("--hidden_path", type=str, default=None)
|
||||
parser.add_argument("--metadata_path", type=str, default=None)
|
||||
parser.add_argument("--export_config_path", type=str, default=None)
|
||||
parser.add_argument("--eval_split", type=str, default=None)
|
||||
parser.add_argument("--dataset_subset_size", type=int, default=None)
|
||||
parser.add_argument("--train_eid_file", type=str, default=None)
|
||||
parser.add_argument("--val_eid_file", type=str, default=None)
|
||||
parser.add_argument("--test_eid_file", type=str, default=None)
|
||||
parser.add_argument("--landmark_start", type=float, default=40.0)
|
||||
parser.add_argument("--landmark_stop", type=float, default=80.0)
|
||||
parser.add_argument(
|
||||
"--landmark_step",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="t_query grid step in years. Default exports every integer age 40..80.",
|
||||
)
|
||||
parser.add_argument("--min_history_events", type=int, default=None)
|
||||
parser.add_argument("--batch_size", type=int, default=None)
|
||||
parser.add_argument("--num_workers", type=int, default=None)
|
||||
parser.add_argument("--device", type=str, default=None)
|
||||
parser.add_argument("--extra_info_types", type=str, default=None)
|
||||
parser.add_argument(
|
||||
"--logits_dtype",
|
||||
type=str,
|
||||
default="float32",
|
||||
choices=["float16", "float32"],
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hidden_dtype",
|
||||
type=str,
|
||||
default="float32",
|
||||
choices=["float16", "float32"],
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
run_path = Path(args.run_path)
|
||||
config_path = run_path / "train_config.json"
|
||||
checkpoint_path = run_path / "best_model.pt"
|
||||
if not config_path.exists():
|
||||
raise FileNotFoundError(f"train_config.json not found: {config_path}")
|
||||
if not checkpoint_path.exists():
|
||||
raise FileNotFoundError(f"best_model.pt not found: {checkpoint_path}")
|
||||
|
||||
cfg = load_json_config(config_path)
|
||||
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
||||
if model_target_mode not in {"next_token", "all_future"}:
|
||||
raise ValueError(f"Unsupported model_target_mode: {model_target_mode!r}")
|
||||
|
||||
target_mode = str(cfg.get("target_mode", "uts"))
|
||||
attn_mask_mode = str(
|
||||
cfg.get(
|
||||
"attn_mask_mode",
|
||||
"non_strict_time" if target_mode == "uts" else "target_aware",
|
||||
)
|
||||
)
|
||||
readout_name = str(
|
||||
cfg.get("readout_name", "same_time_group_end" if target_mode == "uts" else "token")
|
||||
)
|
||||
readout_reduce = str(cfg.get("readout_reduce", "mean"))
|
||||
|
||||
dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset(
|
||||
args,
|
||||
cfg,
|
||||
)
|
||||
validate_dataset_metadata(dataset, cfg)
|
||||
|
||||
landmark_ages = make_landmark_ages(
|
||||
float(args.landmark_start),
|
||||
float(args.landmark_stop),
|
||||
float(args.landmark_step),
|
||||
)
|
||||
first_occurrence_by_token = build_first_occurrence_maps_for_landmarks(
|
||||
dataset,
|
||||
subset_indices,
|
||||
)
|
||||
death_idx = int(dataset.vocab_size) - 1
|
||||
landmark_dataset = LandmarkDataset(
|
||||
dataset=dataset,
|
||||
subset_indices=subset_indices,
|
||||
landmark_ages=landmark_ages,
|
||||
attn_mask_mode=attn_mask_mode,
|
||||
model_target_mode=model_target_mode,
|
||||
min_history_events=int(cfg_get(args, cfg, "min_history_events", 1)),
|
||||
first_occurrence_by_token=first_occurrence_by_token,
|
||||
death_token_ids=[death_idx],
|
||||
)
|
||||
|
||||
state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu")
|
||||
dist_mode = resolve_dist_mode_for_checkpoint(
|
||||
str(cfg.get("dist_mode", "exponential")),
|
||||
state_dict,
|
||||
)
|
||||
cfg_model = dict(cfg)
|
||||
cfg_model["dist_mode"] = dist_mode
|
||||
device = resolve_eval_device(args.device)
|
||||
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
|
||||
load_model_state(model, state_dict)
|
||||
model.eval()
|
||||
|
||||
default_logits_path, default_hidden_path, default_metadata_path, default_config_path = (
|
||||
output_paths_for_run(
|
||||
run_path=run_path,
|
||||
eval_split=eval_split,
|
||||
landmark_start=float(args.landmark_start),
|
||||
landmark_stop=float(args.landmark_stop),
|
||||
landmark_step=float(args.landmark_step),
|
||||
output_dir=args.output_dir,
|
||||
)
|
||||
)
|
||||
logits_path = Path(args.logits_path) if args.logits_path else default_logits_path
|
||||
hidden_path = Path(args.hidden_path) if args.hidden_path else default_hidden_path
|
||||
metadata_path = Path(args.metadata_path) if args.metadata_path else default_metadata_path
|
||||
export_config_path = (
|
||||
Path(args.export_config_path) if args.export_config_path else default_config_path
|
||||
)
|
||||
for path in (logits_path, hidden_path, metadata_path, export_config_path):
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
n_rows = len(landmark_dataset)
|
||||
vocab_size = int(dataset.vocab_size)
|
||||
hidden_dim = int(getattr(model, "n_embd", cfg_get(args, cfg_model, "n_embd", 120)))
|
||||
logits_dtype = numpy_float_dtype(args.logits_dtype)
|
||||
hidden_dtype = numpy_float_dtype(args.hidden_dtype)
|
||||
|
||||
logits_memmap = np.lib.format.open_memmap(
|
||||
logits_path,
|
||||
mode="w+",
|
||||
dtype=logits_dtype,
|
||||
shape=(n_rows, vocab_size),
|
||||
)
|
||||
hidden_memmap = np.lib.format.open_memmap(
|
||||
hidden_path,
|
||||
mode="w+",
|
||||
dtype=hidden_dtype,
|
||||
shape=(n_rows, hidden_dim),
|
||||
)
|
||||
|
||||
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
|
||||
num_workers = int(cfg_get(args, cfg, "num_workers", 4))
|
||||
loader = DataLoader(
|
||||
IndexedLandmarkDataset(landmark_dataset),
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
collate_fn=collate_indexed_landmark_fn,
|
||||
num_workers=num_workers,
|
||||
pin_memory=device.type == "cuda",
|
||||
persistent_workers=num_workers > 0,
|
||||
prefetch_factor=2 if num_workers > 0 else None,
|
||||
)
|
||||
|
||||
print(f"Eval split: {eval_split}")
|
||||
print(f"Split source: {split_source}")
|
||||
print(f"Selected patients: {len(subset_indices)}")
|
||||
print(f"t_query ages: {landmark_ages.tolist()}")
|
||||
print(f"Dist mode: {dist_mode}")
|
||||
print(f"Device: {device}")
|
||||
print(f"Landmark rows: {n_rows}")
|
||||
print(f"Logits: {logits_path} shape={(n_rows, vocab_size)} dtype={logits_dtype}")
|
||||
print(f"Hidden: {hidden_path} shape={(n_rows, hidden_dim)} dtype={hidden_dtype}")
|
||||
print(f"Metadata: {metadata_path}")
|
||||
|
||||
meta_rows: list[dict[str, Any]] = []
|
||||
written = 0
|
||||
with torch.no_grad():
|
||||
for batch in tqdm(loader, desc="Export logits/hidden", dynamic_ncols=True):
|
||||
hidden = infer_landmark_hidden(
|
||||
model=model,
|
||||
batch=batch,
|
||||
device=device,
|
||||
model_target_mode=model_target_mode,
|
||||
readout_name=readout_name,
|
||||
readout_reduce=readout_reduce,
|
||||
)
|
||||
logits = model.calc_risk(hidden)
|
||||
row_indices = batch["row_idx"].detach().cpu().numpy().astype(np.int64)
|
||||
if not np.array_equal(row_indices, np.arange(written, written + len(row_indices))):
|
||||
raise RuntimeError("DataLoader row order changed; export requires shuffle=False.")
|
||||
|
||||
batch_n = int(logits.shape[0])
|
||||
logits_memmap[written : written + batch_n] = (
|
||||
logits.detach().cpu().numpy().astype(logits_dtype, copy=False)
|
||||
)
|
||||
hidden_memmap[written : written + batch_n] = (
|
||||
hidden.detach().cpu().numpy().astype(hidden_dtype, copy=False)
|
||||
)
|
||||
|
||||
for row_idx in row_indices.tolist():
|
||||
meta = landmark_dataset.rows[int(row_idx)]
|
||||
sample = dataset.samples[int(meta["dataset_index"])]
|
||||
meta_rows.append(
|
||||
{
|
||||
"row_index": int(row_idx),
|
||||
"patient_id": int(meta["patient_id"]),
|
||||
"dataset_index": int(meta["dataset_index"]),
|
||||
"eid": int(sample.get("eid", -1)),
|
||||
"sex": int(meta["sex"]),
|
||||
"t_query": float(meta["t_query"]),
|
||||
"landmark_age": float(meta["landmark_age"]),
|
||||
"followup_end_time": float(meta["followup_end_time"]),
|
||||
"death_time": float(meta["death_time"]),
|
||||
}
|
||||
)
|
||||
written += batch_n
|
||||
|
||||
logits_memmap.flush()
|
||||
hidden_memmap.flush()
|
||||
pd.DataFrame(meta_rows).to_csv(metadata_path, index=False)
|
||||
|
||||
export_config = {
|
||||
"run_path": str(run_path),
|
||||
"eval_split": eval_split,
|
||||
"split_source": split_source,
|
||||
"model_target_mode": model_target_mode,
|
||||
"target_mode": target_mode,
|
||||
"attn_mask_mode": attn_mask_mode,
|
||||
"readout_name": readout_name,
|
||||
"readout_reduce": readout_reduce,
|
||||
"dist_mode": dist_mode,
|
||||
"landmark_ages": [float(x) for x in landmark_ages.tolist()],
|
||||
"n_rows": int(n_rows),
|
||||
"vocab_size": int(vocab_size),
|
||||
"hidden_dim": int(hidden_dim),
|
||||
"death_token": int(death_idx),
|
||||
"logits_path": str(logits_path),
|
||||
"hidden_path": str(hidden_path),
|
||||
"metadata_path": str(metadata_path),
|
||||
"logits_dtype": str(logits_dtype),
|
||||
"hidden_dtype": str(hidden_dtype),
|
||||
}
|
||||
with export_config_path.open("w", encoding="utf-8") as f:
|
||||
json.dump(export_config, f, indent=2)
|
||||
|
||||
print(f"Wrote {written} rows.")
|
||||
print(f"Wrote export config: {export_config_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user