JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor AUC evaluation scripts to support model target modes and improve distribution handling

Browse Source
This commit is contained in:
Jiarui Li
2026-06-12 11:33:32 +08:00
parent 0fa8bbbb9a
commit 3125b6119f
3 changed files with 378 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files

123
evaluate_auc.py
View File

@@ -309,6 +309,12 @@ def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: floa
def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth:
model_target_mode = str(cfg_get(
args, cfg, "model_target_mode", "next_token")).lower()
if model_target_mode not in {"next_token", "all_future"}:
raise ValueError(
f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
)
return DeepHealth(
vocab_size=dataset.vocab_size,
n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
@@ -320,7 +326,7 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
n_categories=dataset.n_categories,
cont_type_ids=dataset.cont_type_ids,
n_bins=int(cfg_get(args, cfg, "n_bins", 16)),
target_mode="next_token",
target_mode=model_target_mode,
time_mode=str(cfg_get(args, cfg, "time_mode", "relative")),
dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
@@ -348,15 +354,25 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A
mode = str(cfg_dist_mode).lower()
has_rho_head = any(str(k).startswith("rho_head.")
for k in state_dict.keys())
has_rho_death_head = any(str(k).startswith("rho_death_head.")
for k in state_dict.keys())
if has_rho_head and mode != "weibull":
print(
"[WARN] Checkpoint contains rho_head weights; overriding dist_mode to 'weibull' for evaluation.")
return "weibull"
if has_rho_death_head and mode != "mixed":
print(
"[WARN] Checkpoint contains rho_death_head weights; overriding dist_mode to 'mixed' for evaluation.")
return "mixed"
if (not has_rho_head) and mode == "weibull":
print(
"[WARN] dist_mode is 'weibull' but checkpoint has no rho_head weights; overriding dist_mode to 'exponential'.")
return "exponential"
if (not has_rho_death_head) and mode == "mixed":
print(
"[WARN] dist_mode is 'mixed' but checkpoint has no rho_death_head weights; overriding dist_mode to 'exponential'.")
return "exponential"
return mode
@@ -452,25 +468,27 @@ def infer_readout_hidden(
model: DeepHealth,
loader: DataLoader,
device: torch.device,
model_target_mode: str,
readout_name: str,
readout_reduce: str,
use_amp: bool,
hidden_cache_dtype: str = "float16",
) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
"""
Run the expensive transformer/readout path exactly once and cache hidden states.
"""Cache per-position hidden states used by the unchanged AUC logic."""
model_target_mode = str(model_target_mode).lower()
if model_target_mode not in {"next_token", "all_future"}:
raise ValueError(
f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
)
The older implementation re-ran the whole model once per disease chunk even
though only the final tied risk-head columns changed. That is usually the
dominant bottleneck. This function computes readout hidden states once; later
chunks only perform a cheap selected-vocabulary linear projection.
"""
if readout_name == "same_time_group_end":
readout = None
if model_target_mode == "next_token" and readout_name == "same_time_group_end":
readout = build_readout("same_time_group_end",
reduce=readout_reduce).to(device)
else:
elif model_target_mode == "next_token":
readout = build_readout(readout_name).to(device)
readout.eval()
if readout is not None:
readout.eval()
hidden_parts: List[np.ndarray] = []
arrays: Dict[str, List[np.ndarray]] = {
@@ -501,25 +519,55 @@ def infer_readout_hidden(
if autocast_enabled else contextlib.nullcontext()
)
with amp_context:
hidden = model(
event_seq=event_seq,
time_seq=time_seq,
sex=batch_dev["sex"],
padding_mask=padding_mask,
other_type=batch_dev["other_type"],
other_value=batch_dev["other_value"],
other_value_kind=batch_dev["other_value_kind"],
other_time=batch_dev["other_time"],
target_mode="next_token",
)
ro = readout(
hidden=hidden,
time_seq=time_seq,
padding_mask=padding_mask,
readout_mask=batch_dev["readout_mask"],
)
if model_target_mode == "all_future":
batch_size, seq_len = event_seq.shape
hidden = torch.zeros(
batch_size,
seq_len,
model.n_embd,
device=event_seq.device,
dtype=torch.float32,
)
for pos in range(seq_len):
active = padding_mask[:, pos].bool()
if not active.any():
continue
hidden_pos = model(
event_seq=event_seq[active],
time_seq=time_seq[active],
sex=batch_dev["sex"][active],
padding_mask=padding_mask[active],
t_query=time_seq[active, pos],
other_type=batch_dev["other_type"][active],
other_value=batch_dev["other_value"][active],
other_value_kind=batch_dev["other_value_kind"][active],
other_time=batch_dev["other_time"][active],
target_mode="all_future",
)
hidden[active, pos, :] = hidden_pos.float()
readout_mask_np = batch["padding_mask"].cpu().numpy()
else:
hidden_raw = model(
event_seq=event_seq,
time_seq=time_seq,
sex=batch_dev["sex"],
padding_mask=padding_mask,
other_type=batch_dev["other_type"],
other_value=batch_dev["other_value"],
other_value_kind=batch_dev["other_value_kind"],
other_time=batch_dev["other_time"],
target_mode="next_token",
)
ro = readout(
hidden=hidden_raw,
time_seq=time_seq,
padding_mask=padding_mask,
readout_mask=batch_dev["readout_mask"],
)
hidden = ro.hidden
readout_mask_np = ro.readout_mask.detach().cpu().numpy()
h = ro.hidden.detach().cpu().numpy().astype(out_dtype, copy=False)
h = hidden.detach().cpu().numpy().astype(out_dtype, copy=False)
hidden_parts.append(h)
max_len = max(max_len, h.shape[1])
@@ -529,7 +577,7 @@ def infer_readout_hidden(
batch[k].cpu().numpy().astype(np.int8, copy=False))
else:
arrays[k].append(batch[k].cpu().numpy())
arrays["readout_mask"][-1] = ro.readout_mask.detach().cpu().numpy()
arrays["readout_mask"][-1] = readout_mask_np
def pad_3d(parts: List[np.ndarray], fill: float = 0.0) -> np.ndarray:
out = np.full(
@@ -999,6 +1047,7 @@ def evaluate_auc_pipeline(
age_groups: np.ndarray,
offsets: Sequence[float],
device: torch.device,
model_target_mode: str,
readout_name: str,
readout_reduce: str,
num_workers_auc: int,
@@ -1058,6 +1107,7 @@ def evaluate_auc_pipeline(
model=model,
loader=loader,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
use_amp=use_amp,
@@ -1257,6 +1307,12 @@ def main() -> None:
include_no_event = cfg.get("include_no_event_in_uts_target", False)
target_mode = cfg.get("target_mode", "uts")
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
if model_target_mode not in {"next_token", "all_future"}:
raise ValueError(
"train_config.json model_target_mode must be next_token or all_future, "
f"got {model_target_mode!r}"
)
dist_mode_cfg = cfg.get("dist_mode", "exponential")
readout_name = cfg.get(
"readout_name", "same_time_group_end" if target_mode == "uts" else "token")
@@ -1298,7 +1354,13 @@ def main() -> None:
dist_mode = resolve_dist_mode_for_checkpoint(dist_mode_cfg, state_dict)
cfg = dict(cfg)
cfg["dist_mode"] = dist_mode
cfg["model_target_mode"] = model_target_mode
print(f"Resolved dist_mode for evaluation: {dist_mode}")
print(f"Model target mode for AUC: {model_target_mode}")
print(
"AUC score semantics: evaluate_auc.py uses disease-specific eta/logit scores; "
"dist_mode affects model loading but is not converted to horizon-specific risk probability."
)
model = build_model_from_dataset(args, cfg, dataset).to(device)
load_model_state(model, str(model_ckpt_path),
@@ -1335,6 +1397,7 @@ def main() -> None:
age_groups=age_groups,
offsets=auc_offsets,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
num_workers_auc=int(cfg_get(args, cfg, "num_workers_auc", max(