From 46a3dfe628188dbc87f31336651c64706c83760f Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Sat, 13 Jun 2026 11:42:04 +0800 Subject: [PATCH] Add training scripts for all-future and next-step supervision with DeepHealth - Implement `train_all_future.py` for training with query-conditioned all-future supervision. - Implement `train_next_step.py` for training with next-token/next-time-point supervision. - Introduce `train_util.py` for shared utility functions including logging, dataset splitting, and model checkpointing. - Enhance argument parsing for both training scripts to accommodate new parameters. - Update loss functions and model configurations to support the new training paradigms. --- README.md | 12 +- dataset.py | 97 ++- eval_data.py | 163 +++++ evaluate_auc.py | 12 +- evaluate_auc_v2.py | 6 +- evaluate_doa_auc.py | 547 +++++++++++++++++ evaluate_landmark_auc.py | 7 + evaluate_token_auc.py | 7 + train.py | 1251 -------------------------------------- train_all_future.py | 448 ++++++++++++++ train_next_step.py | 454 ++++++++++++++ train_util.py | 196 ++++++ 12 files changed, 1927 insertions(+), 1273 deletions(-) create mode 100644 eval_data.py create mode 100644 evaluate_doa_auc.py create mode 100644 evaluate_landmark_auc.py create mode 100644 evaluate_token_auc.py delete mode 100644 train.py create mode 100644 train_all_future.py create mode 100644 train_next_step.py create mode 100644 train_util.py diff --git a/README.md b/README.md index 4ffd62c..143942d 100644 --- a/README.md +++ b/README.md @@ -223,7 +223,7 @@ all-future / query-conditioned 监督： ## 训练 -当前 `train.py` 支持 next-token 和 all-future 两类训练入口： +当前 `train_next_step.py` / `train_all_future.py` 支持 next-token 和 all-future 两类训练入口： - `--model_target_mode next_token` - 使用 `NextStepHealthDataset` @@ -235,7 +235,7 @@ all-future / query-conditioned 监督： - 不使用 readout，直接对 query hidden 计算风险 - `--dist_mode exponential/weibull/mixed` 分别搭配 `ExponentialLoss`、`WeibullLoss`、`MixedLoss` -当前 `train.py` 支持所有已有训练目标定义的组合： +当前 `train_next_step.py` / `train_all_future.py` 支持所有已有训练目标定义的组合： | 训练模式 | 时间模式 | 分布/监督 | 默认 loss/readout | | --- | --- | --- | --- | @@ -248,7 +248,7 @@ all-future / query-conditioned 监督： 示例： ```bash -python train.py \ +python train_next_step.py \ --data_prefix ukb \ --labels_file labels.csv \ --model_target_mode next_token \ @@ -262,7 +262,7 @@ python train.py \ all-future 示例： ```bash -python train.py \ +python train_next_step.py \ --data_prefix ukb \ --labels_file labels.csv \ --model_target_mode all_future \ @@ -273,10 +273,10 @@ python train.py \ 选择额外信息变量： ```bash -python train.py --extra_info_types_file extra_info_types_smoking_alcohol_bmi.txt +python train_next_step.py --extra_info_types_file extra_info_types_smoking_alcohol_bmi.txt ``` -`train.py` 只接受 `--extra_info_types_file` 指定变量列表，不接受在 CLI 里直接输入 type id。文件可以每行一个 type id，也可以带 `#` 注释；如果不传 `--extra_info_types_file`，默认使用全部 other-info type。 +`train_next_step.py` / `train_all_future.py` 只接受 `--extra_info_types_file` 指定变量列表，不接受在 CLI 里直接输入 type id。文件可以每行一个 type id，也可以带 `#` 注释；如果不传 `--extra_info_types_file`，默认使用全部 other-info type。 训练输出的 `train_config.json` 会记录： diff --git a/dataset.py b/dataset.py index f26636f..f50df3e 100644 --- a/dataset.py +++ b/dataset.py @@ -93,6 +93,7 @@ def _cache_file_path( split: str | None = None, min_history_events: int | None = None, min_future_events: int | None = None, + validation_query_seed: int | None = None, ) -> str: event_path = f"{data_prefix}_event_data.npy" basic_path = f"{data_prefix}_basic_info.csv" @@ -122,6 +123,7 @@ def _cache_file_path( str(int(include_no_event_in_uts_target)), "" if min_history_events is None else str(int(min_history_events)), "" if min_future_events is None else str(int(min_future_events)), + "" if validation_query_seed is None else str(int(validation_query_seed)), ] for path in (event_path, basic_path, other_path, cate_types_path, labels_file): @@ -305,7 +307,11 @@ class _ExpoBaseDataset(Dataset): "other_time": (rows[:, 4].astype(np.float32) / DAYS_PER_YEAR), } - def _iter_patient_events(self) -> Iterable[tuple[int, np.ndarray, np.ndarray]]: + def _iter_patient_events( + self, + *, + impute_no_event_gaps: bool, + ) -> Iterable[tuple[int, np.ndarray, np.ndarray]]: unique_eids, starts = np.unique(self.event_data[:, 0], return_index=True) ends = np.concatenate([starts[1:], [len(self.event_data)]]) for eid_raw, start, end in zip(unique_eids, starts, ends): @@ -313,7 +319,19 @@ class _ExpoBaseDataset(Dataset): rows = self.event_data[start:end] times_days_raw = rows[:, 1].astype(np.float32) labels_raw = rows[:, 2].astype(np.int64) + + disease_mask = labels_raw != CHECKUP_IDX + times_days_raw = times_days_raw[disease_mask] + labels_raw = labels_raw[disease_mask] + if len(labels_raw) == 0: + yield eid, times_days_raw, labels_raw + continue + labels_raw = np.where(labels_raw >= NO_EVENT_IDX, labels_raw + 1, labels_raw) + if not impute_no_event_gaps: + yield eid, times_days_raw, labels_raw + continue + times_days, labels = _insert_gap_no_event_tokens( times_days_raw, labels_raw, @@ -374,7 +392,7 @@ class NextStepHealthDataset(_ExpoBaseDataset): - UniqueTimeSetExponentialLoss: readout_mask, target_dt_unique, target_multi_hot """ - CACHE_VERSION = 1 + CACHE_VERSION = 2 def __init__( self, @@ -406,7 +424,9 @@ class NextStepHealthDataset(_ExpoBaseDataset): ) self.samples: List[Dict] = [] - for eid, times_days, labels in self._iter_patient_events(): + for eid, times_days, labels in self._iter_patient_events( + impute_no_event_gaps=True, + ): if len(labels) < 2: continue @@ -463,10 +483,12 @@ class AllFutureHealthDataset(_ExpoBaseDataset): targets. Train samples one query time per patient at each __getitem__ call. - Valid/test use deterministic pre-event query points. + Valid/test use random-but-fixed query points. For each patient with N real + disease events, N - 2 query points are sampled from the eligible observed + time range, with at least one future event after every query. """ - CACHE_VERSION = 1 + CACHE_VERSION = 4 def __init__( self, @@ -477,6 +499,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset): include_no_event_in_uts_target: bool = False, min_history_events: int = 1, min_future_events: int = 1, + validation_query_seed: int = 42, extra_info_types: Iterable[int] | None = None, ) -> None: if split not in {"train", "valid", "test"}: @@ -492,6 +515,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset): split=split, min_history_events=min_history_events, min_future_events=min_future_events, + validation_query_seed=validation_query_seed if split in {"valid", "test"} else None, ) cached_state = self._load_cache(cache_path, self.CACHE_VERSION) if cached_state is not None: @@ -509,10 +533,17 @@ class AllFutureHealthDataset(_ExpoBaseDataset): self.split = split self.min_history_events = int(min_history_events) self.min_future_events = int(min_future_events) + self.validation_query_seed = int(validation_query_seed) self.patients: List[Dict] = [] self.valid_queries: List[Tuple[int, float]] = [] + validation_rng = None + if split in {"valid", "test"}: + split_offset = 0 if split == "valid" else 1_000_003 + validation_rng = np.random.RandomState(self.validation_query_seed + split_offset) - for eid, times_days, labels in self._iter_patient_events(): + for eid, times_days, labels in self._iter_patient_events( + impute_no_event_gaps=False, + ): times_years = (times_days / DAYS_PER_YEAR).astype(np.float32) unique_times = np.unique(times_years) if len(labels) < 2 or len(unique_times) < 2: @@ -534,13 +565,18 @@ class AllFutureHealthDataset(_ExpoBaseDataset): self.patients.append(patient) if split in {"valid", "test"}: - for target_time in unique_times[1:]: - t_query = float(target_time) - ONE_DAY_YEARS - if self._is_valid_query(patient, t_query): - self.valid_queries.append((pidx, t_query)) + if validation_rng is None: + raise RuntimeError("validation_rng was not initialized") + self.valid_queries.extend( + (pidx, t_query) + for t_query in self._sample_fixed_validation_queries( + patient, + validation_rng, + ) + ) if split in {"valid", "test"} and not self.valid_queries: - raise ValueError("No valid deterministic query points were built.") + raise ValueError("No random-but-fixed validation query points were built.") self._save_cache(cache_path, self.CACHE_VERSION) @@ -554,6 +590,45 @@ class AllFutureHealthDataset(_ExpoBaseDataset): and patient["t_obs"] > t_query ) + def _sample_fixed_validation_queries( + self, + patient: Dict, + rng: np.random.RandomState, + ) -> List[float]: + times = np.asarray(patient["times"], dtype=np.float32) + labels = np.asarray(patient["labels"], dtype=np.int64) + real_event_mask = ~np.isin( + labels, + np.array([PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64), + ) + real_times = np.sort(times[real_event_mask].astype(np.float32, copy=False)) + n_real_events = int(real_times.size) + n_queries = max(0, n_real_events - 2) + if n_queries == 0: + return [] + + min_hist = int(self.min_history_events) + min_future = int(self.min_future_events) + if n_real_events < min_hist + min_future: + return [] + + left = float(real_times[min_hist - 1]) + right_event_time = float(real_times[n_real_events - min_future]) + right = np.nextafter(np.float32(right_event_time), np.float32(-np.inf)) + if not np.isfinite(left) or not np.isfinite(right) or float(right) <= left: + return [] + + queries: List[float] = [] + max_attempts = max(100, n_queries * 50) + for _ in range(max_attempts): + if len(queries) >= n_queries: + break + t_query = float(rng.uniform(left, float(right))) + if self._is_valid_query(patient, t_query): + queries.append(t_query) + + return queries + def _sample_train_query(self, patient: Dict) -> float: unique_times = np.unique(patient["times"]) if len(unique_times) < 2: diff --git a/eval_data.py b/eval_data.py new file mode 100644 index 0000000..5dc1a51 --- /dev/null +++ b/eval_data.py @@ -0,0 +1,163 @@ +from __future__ import annotations + +from typing import Any, Dict, Iterable, List + +import numpy as np +import torch +from torch.nn.utils.rnn import pad_sequence + +from dataset import AllFutureHealthDataset, HealthDataset +from targets import PAD_IDX + + +class AllFutureSequenceEvalDataset: + """ + Eval-only sequence view for all-future checkpoints. + + All-future training uses pure disease histories, so token-level and landmark + evaluation should not reuse the next-step dataset view that contains + imputed gap tokens. + """ + + def __init__( + self, + data_prefix: str, + labels_file: str, + min_history_events: int = 1, + min_future_events: int = 1, + extra_info_types: Iterable[int] | None = None, + ) -> None: + base = AllFutureHealthDataset( + data_prefix=data_prefix, + labels_file=labels_file, + split="train", + min_history_events=min_history_events, + min_future_events=min_future_events, + extra_info_types=extra_info_types, + ) + + self.base = base + self.label_code_to_id = base.label_code_to_id + self.label_id_to_code = base.label_id_to_code + self.vocab_size = base.vocab_size + self.n_types = base.n_types + self.n_cont_types = base.n_cont_types + self.n_categories = base.n_categories + self.cont_type_ids = base.cont_type_ids + self.extra_info_types = base.extra_info_types + + self.samples: List[Dict[str, Any]] = [] + for patient in base.patients: + labels = np.asarray(patient["labels"], dtype=np.int64) + times = np.asarray(patient["times"], dtype=np.float32) + if labels.size < 2: + continue + input_len = int(labels.size - 1) + self.samples.append( + { + "eid": int(patient["eid"]), + "event_seq": labels[:-1], + "time_seq": times[:-1], + "target_event_seq": labels[1:], + "target_time_seq": times[1:], + "readout_mask": np.ones(input_len, dtype=bool), + "sex": int(patient["sex"]), + "other_type": np.asarray(patient["other_type"], dtype=np.int64), + "other_value": np.asarray(patient["other_value"], dtype=np.float32), + "other_value_kind": np.asarray(patient["other_value_kind"], dtype=np.int64), + "other_time": np.asarray(patient["other_time"], dtype=np.float32), + } + ) + + def __len__(self) -> int: + return len(self.samples) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + s = self.samples[idx] + return { + "event_seq": torch.from_numpy(s["event_seq"]).long(), + "time_seq": torch.from_numpy(s["time_seq"]).float(), + "target_event_seq": torch.from_numpy(s["target_event_seq"]).long(), + "target_time_seq": torch.from_numpy(s["target_time_seq"]).float(), + "readout_mask": torch.from_numpy(s["readout_mask"]).bool(), + "sex": torch.tensor(s["sex"], dtype=torch.long), + "other_type": torch.from_numpy(s["other_type"]).long(), + "other_value": torch.from_numpy(s["other_value"]).float(), + "other_value_kind": torch.from_numpy(s["other_value_kind"]).long(), + "other_time": torch.from_numpy(s["other_time"]).float(), + } + + +def load_sequence_eval_dataset( + *, + model_target_mode: str, + data_prefix: str, + labels_file: str, + no_event_interval_years: float, + include_no_event_in_uts_target: bool, + min_history_events: int, + min_future_events: int, + extra_info_types: Iterable[int] | None, +): + mode = str(model_target_mode).lower() + if mode == "next_token": + return HealthDataset( + data_prefix=data_prefix, + labels_file=labels_file, + no_event_interval_years=no_event_interval_years, + include_no_event_in_uts_target=include_no_event_in_uts_target, + extra_info_types=extra_info_types, + ) + if mode == "all_future": + return AllFutureSequenceEvalDataset( + data_prefix=data_prefix, + labels_file=labels_file, + min_history_events=min_history_events, + min_future_events=min_future_events, + extra_info_types=extra_info_types, + ) + raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}") + + +def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]: + event_seq = pad_sequence( + [s["event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX + ) + time_seq = pad_sequence( + [s["time_seq"] for s in batch], batch_first=True, padding_value=0.0 + ) + target_event_seq = pad_sequence( + [s["target_event_seq"] for s in batch], batch_first=True, padding_value=PAD_IDX + ) + target_time_seq = pad_sequence( + [s["target_time_seq"] for s in batch], batch_first=True, padding_value=0.0 + ) + readout_mask = pad_sequence( + [s["readout_mask"] for s in batch], batch_first=True, padding_value=False + ) + other_type = pad_sequence( + [s["other_type"] for s in batch], batch_first=True, padding_value=0 + ) + other_value = pad_sequence( + [s["other_value"] for s in batch], batch_first=True, padding_value=0.0 + ) + other_value_kind = pad_sequence( + [s["other_value_kind"] for s in batch], batch_first=True, padding_value=0 + ) + other_time = pad_sequence( + [s["other_time"] for s in batch], batch_first=True, padding_value=0.0 + ) + + return { + "event_seq": event_seq, + "time_seq": time_seq, + "padding_mask": event_seq > PAD_IDX, + "target_event_seq": target_event_seq, + "target_time_seq": target_time_seq, + "readout_mask": readout_mask, + "sex": torch.stack([s["sex"] for s in batch]), + "other_type": other_type, + "other_value": other_value, + "other_value_kind": other_value_kind, + "other_time": other_time, + } diff --git a/evaluate_auc.py b/evaluate_auc.py index da867c1..47af00f 100644 --- a/evaluate_auc.py +++ b/evaluate_auc.py @@ -18,7 +18,7 @@ Efficiency notes: avoiding repeated pickling of arrays for every disease. Run from the DeepHealth code directory containing dataset.py, models.py, -readouts.py, and train.py-compatible checkpoints/configs. +readouts.py, and train_config.json-compatible checkpoints/configs. """ from __future__ import annotations @@ -38,7 +38,8 @@ import torch from torch.utils.data import DataLoader, Subset from tqdm.auto import tqdm -from dataset import HealthDataset, collate_fn +from dataset import HealthDataset +from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn from models import DeepHealth from readouts import build_readout from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX @@ -1330,11 +1331,14 @@ def main() -> None: torch.backends.cudnn.benchmark = True print("Loading dataset...") - dataset = HealthDataset( + dataset = load_sequence_eval_dataset( + model_target_mode=model_target_mode, data_prefix=data_prefix, labels_file=labels_file, no_event_interval_years=no_event_interval_years, include_no_event_in_uts_target=include_no_event, + min_history_events=int(cfg.get("all_future_min_history_events", 1)), + min_future_events=int(cfg.get("all_future_min_future_events", 1)), extra_info_types=parse_int_list(cfg.get("extra_info_types", None)), ) validate_dataset_metadata(dataset, cfg) @@ -1346,7 +1350,7 @@ def main() -> None: subset, batch_size=int(cfg_get(args, cfg, "batch_size", 128)), shuffle=False, - collate_fn=collate_fn, + collate_fn=sequence_eval_collate_fn, num_workers=int(cfg_get(args, cfg, "num_workers", 4)), pin_memory=device.type == "cuda", persistent_workers=int(cfg_get(args, cfg, "num_workers", 4)) > 0, diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py index 8aef743..bbd922b 100644 --- a/evaluate_auc_v2.py +++ b/evaluate_auc_v2.py @@ -28,6 +28,7 @@ from torch.utils.data import DataLoader, Dataset from tqdm.auto import tqdm from dataset import HealthDataset +from eval_data import load_sequence_eval_dataset from models import DeepHealth from readouts import build_readout from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX @@ -1391,11 +1392,14 @@ def main() -> None: labels_meta = pd.read_csv(str(labels_meta_path)) print("Loading dataset...") - dataset = HealthDataset( + dataset = load_sequence_eval_dataset( + model_target_mode=model_target_mode, data_prefix=data_prefix, labels_file=labels_file, no_event_interval_years=float(no_event_interval_years), include_no_event_in_uts_target=bool(include_no_event_in_uts_target), + min_history_events=int(cfg.get("all_future_min_history_events", 1)), + min_future_events=int(cfg.get("all_future_min_future_events", 1)), extra_info_types=parse_int_list(cfg.get("extra_info_types", None)), ) validate_dataset_metadata(dataset, cfg) diff --git a/evaluate_doa_auc.py b/evaluate_doa_auc.py new file mode 100644 index 0000000..f30b7e8 --- /dev/null +++ b/evaluate_doa_auc.py @@ -0,0 +1,547 @@ +"""Evaluate disease AUC at date of assessment (DOA). + +Cases are patients whose first occurrence of a disease is after DOA and within +the requested horizon. Controls are patients who never have that disease in the +full observed record. Patients prevalent at/before DOA or incident after the +horizon are not used for that disease-horizon AUC. + +The script adapts automatically to checkpoint target mode: + - next_token: use the DOA token position, inserting at DOA when no + real disease token exists at DOA; + - all_future: query the model directly with t_query=DOA, allowing empty + disease history because other-info tokens still describe the DOA state. +""" +from __future__ import annotations + +import argparse +import contextlib +import json +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple + +import numpy as np +import pandas as pd +import torch +from torch.nn.utils.rnn import pad_sequence +from torch.utils.data import DataLoader, Dataset +from tqdm.auto import tqdm + +from dataset import _ExpoBaseDataset +from evaluate_auc_v2 import ( + build_metadata_for_merge, + build_model_from_dataset, + get_auc_delong_var, + load_checkpoint_state_dict, + load_json_config, + load_model_state, + parse_float_list, + parse_int_list, + project_distribution_chunk, + resolve_dist_mode_for_checkpoint, + select_disease_tokens, + validate_dataset_metadata, + _score_to_probability, +) +from readouts import build_readout +from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX, DAYS_PER_YEAR + + +SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX} + + +def cfg_get(args: argparse.Namespace, cfg: Dict[str, Any], name: str, default: Any) -> Any: + value = getattr(args, name, None) + if value is not None: + return value + return cfg.get(name, default) + + +class DOAStatusDataset(_ExpoBaseDataset): + def __init__( + self, + data_prefix: str, + labels_file: str, + model_target_mode: str, + extra_info_types: Iterable[int] | None = None, + ) -> None: + super().__init__( + data_prefix=data_prefix, + labels_file=labels_file, + no_event_interval_years=5.0, + include_no_event_in_uts_target=False, + extra_info_types=extra_info_types, + ) + self.model_target_mode = str(model_target_mode).lower() + if self.model_target_mode not in {"next_token", "all_future"}: + raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}") + + self.records: List[Dict[str, Any]] = [] + self.first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]] = {} + + unique_eids, starts = np.unique(self.event_data[:, 0], return_index=True) + ends = np.concatenate([starts[1:], [len(self.event_data)]]) + first_lists: Dict[int, List[Tuple[int, float]]] = {} + + for eid_raw, start, end in zip(unique_eids, starts, ends): + eid = int(eid_raw) + rows = self.event_data[start:end] + checkup_rows = rows[rows[:, 2].astype(np.int64) == CHECKUP_IDX] + if len(checkup_rows) == 0: + continue + + features = self._split_features(eid) + if features is None: + continue + + doa_days = float(np.min(checkup_rows[:, 1].astype(np.float32))) + doa_years = np.float32(doa_days / DAYS_PER_YEAR) + + disease_rows = rows[rows[:, 2].astype(np.int64) != CHECKUP_IDX] + disease_times = disease_rows[:, 1].astype(np.float32) / DAYS_PER_YEAR + disease_labels_raw = disease_rows[:, 2].astype(np.int64) + disease_labels = np.where( + disease_labels_raw >= NO_EVENT_IDX, + disease_labels_raw + 1, + disease_labels_raw, + ).astype(np.int64) + order = np.lexsort((disease_labels, disease_times)) + disease_times = disease_times[order].astype(np.float32) + disease_labels = disease_labels[order].astype(np.int64) + + patient_id = len(self.records) + for token in np.unique(disease_labels).tolist(): + token = int(token) + if token in SPECIAL_TOKENS: + continue + hit = np.where(disease_labels == token)[0] + if hit.size: + first_lists.setdefault(token, []).append( + (patient_id, float(disease_times[int(hit[0])])) + ) + + hist = disease_times <= doa_years + hist_events = disease_labels[hist] + hist_times = disease_times[hist] + + if self.model_target_mode == "next_token": + at_doa = np.isclose(hist_times, doa_years, rtol=0.0, atol=1e-6) + if hist_events.size == 0 or not np.any(at_doa): + event_seq = np.concatenate([ + hist_events, + np.array([NO_EVENT_IDX], dtype=np.int64), + ]) + time_seq = np.concatenate([ + hist_times, + np.array([doa_years], dtype=np.float32), + ]) + else: + event_seq = hist_events + time_seq = hist_times + readout_pos = int(len(event_seq) - 1) + else: + event_seq = hist_events + time_seq = hist_times + readout_pos = -1 + + self.records.append( + { + "patient_id": patient_id, + "eid": eid, + "doa": doa_years, + "event_seq": event_seq.astype(np.int64), + "time_seq": time_seq.astype(np.float32), + "readout_pos": readout_pos, + "full_events": disease_labels, + "full_times": disease_times, + "sex": int(features["sex"]), + "other_type": features["other_type"], + "other_value": features["other_value"], + "other_value_kind": features["other_value_kind"], + "other_time": features["other_time"], + } + ) + + for token, pairs in first_lists.items(): + self.first_occurrence_by_token[int(token)] = ( + np.asarray([p for p, _ in pairs], dtype=np.int32), + np.asarray([t for _, t in pairs], dtype=np.float32), + ) + + if not self.records: + raise RuntimeError("No DOA records were built from checkup events.") + + def __len__(self) -> int: + return len(self.records) + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + s = self.records[idx] + return { + "event_seq": torch.from_numpy(s["event_seq"]).long(), + "time_seq": torch.from_numpy(s["time_seq"]).float(), + "readout_pos": torch.tensor(s["readout_pos"], dtype=torch.long), + "t_query": torch.tensor(float(s["doa"]), dtype=torch.float32), + "patient_id": torch.tensor(s["patient_id"], dtype=torch.long), + "sex": torch.tensor(s["sex"], dtype=torch.long), + "other_type": torch.from_numpy(s["other_type"]).long(), + "other_value": torch.from_numpy(s["other_value"]).float(), + "other_value_kind": torch.from_numpy(s["other_value_kind"]).long(), + "other_time": torch.from_numpy(s["other_time"]).float(), + } + + +def collate_doa_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]: + event_seq = pad_sequence( + [x["event_seq"] for x in batch], batch_first=True, padding_value=PAD_IDX + ) + time_seq = pad_sequence( + [x["time_seq"] for x in batch], batch_first=True, padding_value=0.0 + ) + other_type = pad_sequence( + [x["other_type"] for x in batch], batch_first=True, padding_value=0 + ) + other_value = pad_sequence( + [x["other_value"] for x in batch], batch_first=True, padding_value=0.0 + ) + other_value_kind = pad_sequence( + [x["other_value_kind"] for x in batch], batch_first=True, padding_value=0 + ) + other_time = pad_sequence( + [x["other_time"] for x in batch], batch_first=True, padding_value=0.0 + ) + + readout_mask = torch.zeros_like(event_seq, dtype=torch.bool) + readout_pos = torch.stack([x["readout_pos"] for x in batch]) + for i, pos in enumerate(readout_pos.tolist()): + if pos >= 0: + readout_mask[i, int(pos)] = True + + return { + "event_seq": event_seq, + "time_seq": time_seq, + "padding_mask": event_seq > PAD_IDX, + "readout_mask": readout_mask, + "readout_pos": readout_pos, + "t_query": torch.stack([x["t_query"] for x in batch]), + "patient_id": torch.stack([x["patient_id"] for x in batch]), + "sex": torch.stack([x["sex"] for x in batch]), + "other_type": other_type, + "other_value": other_value, + "other_value_kind": other_value_kind, + "other_time": other_time, + } + + +@torch.inference_mode() +def infer_doa_hidden( + model, + loader: DataLoader, + device: torch.device, + model_target_mode: str, + readout_name: str, + readout_reduce: str, + use_amp: bool, +) -> Tuple[np.ndarray, Dict[str, np.ndarray]]: + model_target_mode = str(model_target_mode).lower() + readout = None + if model_target_mode == "next_token": + if readout_name == "same_time_group_end": + readout = build_readout("same_time_group_end", reduce=readout_reduce).to(device) + else: + readout = build_readout(readout_name).to(device) + readout.eval() + + hidden_parts: List[np.ndarray] = [] + patient_parts: List[np.ndarray] = [] + sex_parts: List[np.ndarray] = [] + autocast_enabled = bool(use_amp and device.type == "cuda") + + for batch in tqdm(loader, desc="DOA inference", leave=False, dynamic_ncols=True): + batch_dev = { + k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v) + for k, v in batch.items() + } + amp_context = ( + torch.autocast(device_type=device.type, dtype=torch.float16) + if autocast_enabled + else contextlib.nullcontext() + ) + with amp_context: + if model_target_mode == "all_future": + hidden = model( + event_seq=batch_dev["event_seq"], + time_seq=batch_dev["time_seq"], + sex=batch_dev["sex"], + padding_mask=batch_dev["padding_mask"], + t_query=batch_dev["t_query"], + 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="all_future", + ) + else: + hidden_raw = model( + event_seq=batch_dev["event_seq"], + time_seq=batch_dev["time_seq"], + sex=batch_dev["sex"], + padding_mask=batch_dev["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=batch_dev["time_seq"], + padding_mask=batch_dev["padding_mask"], + readout_mask=batch_dev["readout_mask"], + ) + if ro.hidden.dim() == 2: + hidden = ro.hidden + else: + hidden = ro.hidden[batch_dev["readout_mask"]] + + hidden_parts.append(hidden.detach().float().cpu().numpy().astype(np.float32, copy=False)) + patient_parts.append(batch["patient_id"].cpu().numpy().astype(np.int32, copy=False)) + sex_parts.append(batch["sex"].cpu().numpy().astype(np.int8, copy=False)) + + return ( + np.concatenate(hidden_parts, axis=0), + { + "patient_id": np.concatenate(patient_parts, axis=0), + "sex": np.concatenate(sex_parts, axis=0), + }, + ) + + +def first_time_array( + first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]], + token: int, + patient_count: int, +) -> np.ndarray: + out = np.full(patient_count, np.inf, dtype=np.float32) + pairs = first_occurrence_by_token.get(int(token)) + if pairs is not None: + p, t = pairs + out[np.asarray(p, dtype=np.int64)] = np.asarray(t, dtype=np.float32) + return out + + +def evaluate_doa_auc( + dataset: DOAStatusDataset, + hidden_all: np.ndarray, + row_arrays: Dict[str, np.ndarray], + model, + disease_ids: Sequence[int], + horizons: np.ndarray, + dist_mode: str, + score_mode: str, + min_cases: int, + device: torch.device, + logit_batch_size: int, + use_amp: bool, +) -> pd.DataFrame: + logits_all, rho_all = project_distribution_chunk( + model=model, + hidden_all=hidden_all, + disease_ids=disease_ids, + dist_mode=dist_mode, + device=device, + logit_batch_size=logit_batch_size, + use_amp=use_amp, + ) + patient_ids = row_arrays["patient_id"].astype(np.int32) + sex = row_arrays["sex"].astype(np.int8) + doa = np.asarray([r["doa"] for r in dataset.records], dtype=np.float32)[patient_ids] + patient_count = len(dataset.records) + death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", 0) - 1)) + + rows: List[Dict[str, Any]] = [] + for col, token in enumerate([int(x) for x in disease_ids]): + first_time = first_time_array(dataset.first_occurrence_by_token, token, patient_count)[patient_ids] + never = np.isinf(first_time) + incident_after_doa = first_time > doa + + for horizon in horizons.tolist(): + horizon = float(horizon) + case_mask = incident_after_doa & (first_time <= doa + np.float32(horizon)) + control_mask = never + if int(case_mask.sum()) < min_cases or int(control_mask.sum()) < min_cases: + continue + + rho_col = None if rho_all is None else rho_all[:, col] + scores = _score_to_probability( + logits=logits_all[:, col], + rho=rho_col, + score_mode=score_mode, + horizon=horizon, + dist_mode=dist_mode, + token=token, + death_idx=death_idx, + ) + + for sex_value, sex_name in [(0, "female"), (1, "male"), (-1, "all")]: + if sex_value == -1: + sex_mask = np.ones_like(case_mask, dtype=bool) + else: + sex_mask = sex == sex_value + cm = case_mask & sex_mask + nm = control_mask & sex_mask + if int(cm.sum()) < min_cases or int(nm.sum()) < min_cases: + continue + auc, var = get_auc_delong_var(scores[cm], scores[nm]) + rows.append( + { + "token": token, + "horizon": horizon, + "sex": sex_name, + "n_case": int(cm.sum()), + "n_control": int(nm.sum()), + "auc": auc, + "auc_var": var, + "auc_se": float(np.sqrt(max(var, 0.0))) if np.isfinite(var) else np.nan, + } + ) + return pd.DataFrame(rows) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Evaluate DOA fixed-horizon disease AUC") + parser.add_argument("--run_path", type=str, required=True) + parser.add_argument("--output_path", type=str, default=None) + parser.add_argument("--batch_size", type=int, default=None) + parser.add_argument("--num_workers", type=int, default=None) + parser.add_argument("--logit_batch_size", type=int, default=None) + parser.add_argument("--horizons", type=str, default=None) + parser.add_argument("--score_mode", type=str, choices=["risk", "eta"], default=None) + parser.add_argument("--filter_min_total", type=int, default=None) + parser.add_argument("--min_cases", type=int, default=None) + parser.add_argument("--labels_meta_path", type=str, default=None) + parser.add_argument("--use_amp", action=argparse.BooleanOptionalAction, default=None) + args = parser.parse_args() + + run_path = Path(args.run_path) + cfg = load_json_config(run_path / "train_config.json") + ckpt_path = run_path / "best_model.pt" + if not ckpt_path.exists(): + raise FileNotFoundError(f"best_model.pt not found in {run_path}") + + output_path = Path(args.output_path or run_path) + output_path.mkdir(parents=True, exist_ok=True) + + 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}") + + labels_meta_path = cfg_get(args, cfg, "labels_meta_path", None) + if labels_meta_path is None: + labels_meta_path = cfg.get("labels_meta_path", "delphi_labels_chapters_colours_icd.csv") + labels_meta = pd.read_csv(labels_meta_path) if labels_meta_path and Path(labels_meta_path).exists() else None + + dataset = DOAStatusDataset( + data_prefix=cfg.get("data_prefix", "ukb"), + labels_file=cfg.get("labels_file", "labels.csv"), + model_target_mode=model_target_mode, + extra_info_types=parse_int_list(cfg.get("extra_info_types", None)), + ) + validate_dataset_metadata(dataset, cfg) + + disease_requested = parse_int_list(cfg_get(args, cfg, "diseases_of_interest", None)) + disease_ids = select_disease_tokens( + dataset=dataset, + labels_meta=labels_meta, + requested_tokens=disease_requested, + filter_min_total=int(cfg_get(args, cfg, "filter_min_total", 0)), + first_occurrence_by_token=dataset.first_occurrence_by_token, + ) + if not disease_ids: + raise RuntimeError("No disease tokens selected after filtering.") + + horizons = np.asarray( + parse_float_list(cfg_get(args, cfg, "horizons", "1,5,10")) or [1.0, 5.0, 10.0], + dtype=np.float32, + ) + score_mode = str(cfg_get(args, cfg, "score_mode", "risk")).lower() + min_cases = int(cfg_get(args, cfg, "min_cases", 2)) + + state_dict = load_checkpoint_state_dict(ckpt_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 = torch.device(cfg.get("device", "cuda") if torch.cuda.is_available() else "cpu") + model = build_model_from_dataset(args, cfg_model, dataset).to(device) + load_model_state(model, state_dict) + model.eval() + + if model_target_mode == "next_token" and ( + model.token_embedding.num_embeddings <= NO_EVENT_IDX + or model.risk_head.out_features <= NO_EVENT_IDX + ): + raise RuntimeError("Next-token DOA evaluation requires in the model vocabulary.") + + loader = DataLoader( + dataset, + batch_size=int(cfg_get(args, cfg, "batch_size", 128)), + shuffle=False, + collate_fn=collate_doa_fn, + num_workers=int(cfg_get(args, cfg, "num_workers", 4)), + pin_memory=device.type == "cuda", + persistent_workers=int(cfg_get(args, cfg, "num_workers", 4)) > 0, + prefetch_factor=2 if int(cfg_get(args, cfg, "num_workers", 4)) > 0 else None, + ) + + target_mode = cfg.get("target_mode", "uts") + 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")) + + print(f"DOA records: {len(dataset)}") + print(f"Model target mode: {model_target_mode}") + print(f"Dist mode: {dist_mode}") + print(f"Score mode: {score_mode}") + print(f"Horizons: {horizons.tolist()}") + print(f"Disease tokens: {len(disease_ids)}") + + hidden_all, row_arrays = infer_doa_hidden( + model=model, + loader=loader, + device=device, + model_target_mode=model_target_mode, + readout_name=readout_name, + readout_reduce=readout_reduce, + use_amp=bool(cfg_get(args, cfg, "use_amp", False)), + ) + result = evaluate_doa_auc( + dataset=dataset, + hidden_all=hidden_all, + row_arrays=row_arrays, + model=model, + disease_ids=disease_ids, + horizons=horizons, + dist_mode=dist_mode, + score_mode=score_mode, + min_cases=min_cases, + device=device, + logit_batch_size=int(cfg_get(args, cfg, "logit_batch_size", cfg_get(args, cfg, "batch_size", 128))), + use_amp=bool(cfg_get(args, cfg, "use_amp", False)), + ) + if result.empty: + raise RuntimeError("No DOA AUC rows produced. Check disease selection and min_cases.") + + meta = build_metadata_for_merge(dataset, labels_meta) + result = result.merge(meta, on="token", how="left") + out_file = output_path / "doa_auc.csv" + result.to_csv(out_file, index=False) + + summary = result.groupby(["token", "label_code", "horizon"], dropna=False, as_index=False).agg( + auc_mean=("auc", "mean"), + n_case=("n_case", "sum"), + n_control=("n_control", "sum"), + ) + summary.to_csv(output_path / "doa_auc_summary.csv", index=False) + print(f"Wrote {out_file}") + + +if __name__ == "__main__": + main() diff --git a/evaluate_landmark_auc.py b/evaluate_landmark_auc.py new file mode 100644 index 0000000..25b0ad4 --- /dev/null +++ b/evaluate_landmark_auc.py @@ -0,0 +1,7 @@ +from __future__ import annotations + +from evaluate_auc_v2 import main + + +if __name__ == "__main__": + main() diff --git a/evaluate_token_auc.py b/evaluate_token_auc.py new file mode 100644 index 0000000..17d6330 --- /dev/null +++ b/evaluate_token_auc.py @@ -0,0 +1,7 @@ +from __future__ import annotations + +from evaluate_auc import main + + +if __name__ == "__main__": + main() diff --git a/train.py b/train.py deleted file mode 100644 index c0501ab..0000000 --- a/train.py +++ /dev/null @@ -1,1251 +0,0 @@ -""" -Training script for DeepHealth model. - -Implements the complete training pipeline: -1. Load data and split train/val/test -2. Build model, optimizer, readout, and loss -3. Train with adaptive learning rate (warmup + cosine annealing) -4. Early stopping based on validation loss -5. Save checkpoints and metrics -""" - -from __future__ import annotations - -import argparse -import json -import logging -import math -import os -import sys -import time -from datetime import datetime -from pathlib import Path -from typing import Any, Dict, Tuple - -import numpy as np -import torch -import torch.nn as nn -from torch.utils.data import DataLoader, RandomSampler, Subset -from torch.optim import AdamW -from torch.nn.utils import clip_grad_norm_ -from tqdm.auto import tqdm - -from dataset import ( - AllFutureHealthDataset, - HealthDataset, - all_future_collate_fn, - collate_fn, -) -from models import DeepHealth -from readouts import build_readout -from losses import build_loss -from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX - - -# --------------------------------------------------------------------------- -# Setup Logging -# --------------------------------------------------------------------------- - -def setup_logging(run_dir: Path) -> logging.Logger: - """Configure logging to both console and file.""" - run_dir.mkdir(parents=True, exist_ok=True) - log_file = run_dir / "train.log" - - logger = logging.getLogger("DeepHealth") - logger.setLevel(logging.INFO) - logger.handlers.clear() - - # Console handler - console_handler = logging.StreamHandler(sys.stdout) - console_handler.setLevel(logging.INFO) - console_formatter = logging.Formatter( - "%(asctime)s - %(levelname)s - %(message)s", - datefmt="%Y-%m-%d %H:%M:%S" - ) - console_handler.setFormatter(console_formatter) - logger.addHandler(console_handler) - - # File handler - file_handler = logging.FileHandler(log_file, mode="w") - file_handler.setLevel(logging.INFO) - file_formatter = logging.Formatter( - "%(asctime)s - %(levelname)s - %(message)s", - datefmt="%Y-%m-%d %H:%M:%S" - ) - file_handler.setFormatter(file_formatter) - logger.addHandler(file_handler) - - return logger - - -# --------------------------------------------------------------------------- -# Utilities -# --------------------------------------------------------------------------- - -def set_seed(seed: int) -> None: - """Set random seed for reproducibility.""" - np.random.seed(seed) - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - - -def load_extra_info_types_file(path: str) -> list[int]: - """ - Load other-information type ids from a text or JSON file. - - Text files may use whitespace, commas, or semicolons as separators. Lines may - contain comments after '#'. JSON files should contain a top-level list of ids. - """ - file_path = Path(path) - if not file_path.exists(): - raise FileNotFoundError(f"extra_info_types_file not found: {path}") - if not file_path.is_file(): - raise ValueError(f"extra_info_types_file is not a file: {path}") - - text = file_path.read_text(encoding="utf-8").strip() - if not text: - return [] - - if text.startswith("["): - try: - raw_items = json.loads(text) - except json.JSONDecodeError as exc: - raise ValueError( - f"Invalid JSON in extra_info_types_file: {path}" - ) from exc - if not isinstance(raw_items, list): - raise ValueError( - f"extra_info_types_file JSON must be a list, got {type(raw_items).__name__}" - ) - else: - tokens: list[str] = [] - for line in text.splitlines(): - line = line.split("#", 1)[0].strip() - if not line: - continue - tokens.extend(line.replace(",", " ").replace(";", " ").split()) - raw_items = tokens - - parsed: list[int] = [] - for item in raw_items: - try: - type_id = int(item) - except (TypeError, ValueError) as exc: - raise ValueError( - f"Invalid extra info type id {item!r} in {path}; expected integers." - ) from exc - parsed.append(type_id) - - return parsed - - -def configure_torch_for_training(device: torch.device) -> None: - """Enable backend settings that can improve training throughput on CUDA.""" - if device.type == "cuda": - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - if hasattr(torch, "set_float32_matmul_precision"): - torch.set_float32_matmul_precision("high") - - -def resolve_device(device_arg: str) -> torch.device: - """Resolve and validate requested device string.""" - requested = device_arg.strip().lower() - - if requested == "cpu": - return torch.device("cpu") - - if requested == "cuda": - if not torch.cuda.is_available(): - return torch.device("cpu") - return torch.device("cuda") - - if requested.startswith("cuda:"): - if not torch.cuda.is_available(): - return torch.device("cpu") - try: - index = int(requested.split(":", 1)[1]) - except ValueError as exc: - raise ValueError( - f"Invalid CUDA device format '{device_arg}'. Use 'cuda' or 'cuda:'." - ) from exc - - n_cuda = torch.cuda.device_count() - if index < 0 or index >= n_cuda: - raise ValueError( - f"Requested device '{device_arg}' is out of range. " - f"Available CUDA devices: 0..{max(0, n_cuda - 1)}" - ) - return torch.device(f"cuda:{index}") - - raise ValueError( - f"Unsupported device '{device_arg}'. Use 'cpu', 'cuda', or 'cuda:'." - ) - - -def split_dataset( - dataset: HealthDataset, - train_ratio: float, - val_ratio: float, - test_ratio: float, - seed: int, -) -> Tuple[Subset, Subset, Subset]: - """ - Split dataset into train/val/test subsets. - - Parameters - ---------- - train_ratio, val_ratio, test_ratio : float - Ratios must sum to 1.0 (within 1e-6 tolerance). - seed : int - Random seed for splitting. - - Returns - ------- - train_subset, val_subset, test_subset - """ - total = train_ratio + val_ratio + test_ratio - if not np.isclose(total, 1.0, atol=1e-6): - raise ValueError( - f"train_ratio + val_ratio + test_ratio must equal 1.0, " - f"got {total}" - ) - - n = len(dataset) - rng = np.random.RandomState(seed) - indices = rng.permutation(n) - - n_train = int(n * train_ratio) - n_val = int(n * val_ratio) - - train_indices = indices[:n_train] - val_indices = indices[n_train: n_train + n_val] - test_indices = indices[n_train + n_val:] - - return ( - Subset(dataset, train_indices), - Subset(dataset, val_indices), - Subset(dataset, test_indices), - ) - - -def split_all_future_datasets( - train_dataset: AllFutureHealthDataset, - val_dataset: AllFutureHealthDataset, - test_dataset: AllFutureHealthDataset, - train_ratio: float, - val_ratio: float, - test_ratio: float, - seed: int, -) -> Tuple[Subset, Subset, Subset]: - """Split all-future datasets by patient, then select validation/test queries.""" - total = train_ratio + val_ratio + test_ratio - if not np.isclose(total, 1.0, atol=1e-6): - raise ValueError( - f"train_ratio + val_ratio + test_ratio must equal 1.0, got {total}" - ) - - n_patients = len(train_dataset.patients) - rng = np.random.RandomState(seed) - patient_indices = rng.permutation(n_patients) - n_train = int(n_patients * train_ratio) - n_val = int(n_patients * val_ratio) - - train_patient_idx = patient_indices[:n_train] - val_patient_set = set(int(x) for x in patient_indices[n_train:n_train + n_val]) - test_patient_set = set(int(x) for x in patient_indices[n_train + n_val:]) - - val_query_idx = [ - i for i, (pidx, _t_query) in enumerate(val_dataset.valid_queries) - if int(pidx) in val_patient_set - ] - test_query_idx = [ - i for i, (pidx, _t_query) in enumerate(test_dataset.valid_queries) - if int(pidx) in test_patient_set - ] - - if not val_query_idx: - raise ValueError("All-future validation split has no valid query samples.") - if not test_query_idx: - raise ValueError("All-future test split has no valid query samples.") - - return ( - Subset(train_dataset, train_patient_idx), - Subset(val_dataset, np.asarray(val_query_idx, dtype=np.int64)), - Subset(test_dataset, np.asarray(test_query_idx, dtype=np.int64)), - ) - - -def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth: - """ - Build DeepHealth model using metadata from dataset. - - Uses unified other-information metadata computed during dataset initialization. - """ - model = DeepHealth( - vocab_size=dataset.vocab_size, - n_embd=args.n_embd, - n_head=args.n_head, - n_hist_layer=args.n_hist_layer, - n_tab_layer=args.n_tab_layer, - n_types=dataset.n_types, - n_cont_types=dataset.n_cont_types, - n_categories=dataset.n_categories, - cont_type_ids=dataset.cont_type_ids, - n_bins=args.n_bins, - target_mode=args.model_target_mode, - time_mode=args.time_mode, - dist_mode=args.dist_mode, - dropout=args.dropout, - ) - - return model - - -def build_optimizer(args: argparse.Namespace, model: nn.Module) -> AdamW: - """Build AdamW optimizer.""" - return AdamW( - model.parameters(), - lr=args.base_lr, - betas=args.betas, - weight_decay=args.weight_decay, - ) - - -def get_lr( - epoch: int, - args: argparse.Namespace, - adaptive_lr: float, - total_epochs: int, -) -> float: - """ - Calculate learning rate for the given epoch. - - Warmup: linear from 0 to adaptive_lr over warmup_epochs - After: cosine annealing from adaptive_lr to adaptive_lr * min_lr_ratio - """ - if epoch < args.warmup_epochs: - # Linear warmup - return adaptive_lr * (epoch + 1) / args.warmup_epochs - else: - # Cosine annealing - progress = (epoch - args.warmup_epochs) / \ - (total_epochs - args.warmup_epochs) - return adaptive_lr * (args.min_lr_ratio + 0.5 * (1 + math.cos(math.pi * progress)) * (1 - args.min_lr_ratio)) - - -def set_optimizer_lr(optimizer: AdamW, lr: float) -> None: - """Set learning rate for all parameter groups.""" - for param_group in optimizer.param_groups: - param_group["lr"] = lr - - -def move_batch_to_device(batch: Dict, device: torch.device) -> Dict: - """Move all tensors in batch to specified device.""" - non_blocking = device.type == "cuda" - return { - key: value.to(device, non_blocking=non_blocking) if isinstance( - value, torch.Tensor) else value - for key, value in batch.items() - } - - -def compute_loss( - args: argparse.Namespace, - model: DeepHealth, - readout: nn.Module | None, - criterion: nn.Module, - batch: Dict[str, torch.Tensor], - device: torch.device, -) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]: - """ - Compute loss for one batch. - - Flow: forward model, apply readout, compute risk logits, compute loss. - - Parameters - ---------- - args : argparse.Namespace - Training configuration with target_mode and loss options. - model : DeepHealth - The model. - readout : nn.Module - Readout module. - criterion : nn.Module - Loss criterion. - batch : Dict - Batch data from DataLoader. - device : torch.device - Device to compute on. - - Returns - ------- - loss : torch.Tensor - Scalar loss tensor. - """ - # Move batch to device - batch = move_batch_to_device(batch, device) - - event_seq = batch["event_seq"] # (B, L) - time_seq = batch["time_seq"] # (B, L) - padding_mask = batch["padding_mask"] # (B, L) - sex = batch["sex"] # (B,) - - if args.model_target_mode == "all_future": - hidden = model( - event_seq=event_seq, - time_seq=time_seq, - sex=sex, - padding_mask=padding_mask, - t_query=batch["t_query"], - other_type=batch["other_type"], - other_value=batch["other_value"], - other_value_kind=batch["other_value_kind"], - other_time=batch["other_time"], - target_mode="all_future", - ) - logits = model.calc_risk(hidden) - - if args.dist_mode == "exponential": - loss = criterion( - logits=logits, - targets=batch["future_targets"], - exposure=batch["exposure"], - ) - elif args.dist_mode == "weibull": - loss = criterion( - logits=logits, - weibull_rho=model.calc_weibull_rho(hidden), - targets=batch["future_targets"], - dt=batch["future_dt"], - exposure=batch["exposure"], - ) - elif args.dist_mode == "mixed": - loss = criterion( - logits=logits, - death_rho=model.calc_death_rho(hidden), - targets=batch["future_targets"], - dt=batch["future_dt"], - exposure=batch["exposure"], - ) - else: - raise ValueError(f"Unknown dist_mode: {args.dist_mode}") - loss_parts = {"total": loss.detach()} - else: - if readout is None: - raise ValueError("next_token training requires a readout module") - - hidden = model( - event_seq=event_seq, - time_seq=time_seq, - sex=sex, - padding_mask=padding_mask, - other_type=batch["other_type"], - other_value=batch["other_value"], - other_value_kind=batch["other_value_kind"], - other_time=batch["other_time"], - target_mode="next_token", - ) - - # Apply readout - readout_mask = ( - batch["readout_mask"] - if args.readout_name == "same_time_group_end" - else None - ) - readout_out = readout( - hidden=hidden, - time_seq=time_seq, - padding_mask=padding_mask, - readout_mask=readout_mask, - ) - - # Compute risk logits - logits = model.calc_risk(readout_out.hidden) - - # Compute loss based on target_mode - if args.target_mode == "delphi2m": - loss_out = criterion( - logits=logits, - target_events=batch["target_event_seq"], - target_times=batch["target_time_seq"], - current_times=batch["time_seq"], - padding_mask=readout_out.readout_mask, - return_components=True, - ) - elif args.target_mode == "uts": - loss_out = criterion( - logits=logits, - target_multi_hot=batch["target_multi_hot"], - target_dt_unique=batch["target_dt_unique"], - readout_mask=readout_out.readout_mask, - return_components=True, - ) - else: - raise ValueError(f"Unknown target_mode: {args.target_mode}") - - loss, loss_parts = loss_out - - # Check for NaN/Inf - if not torch.isfinite(loss): - raise RuntimeError( - f"Loss is not finite: {loss.item()}. " - f"batch_idx info: event_seq shape {event_seq.shape}, " - f"logits shape {logits.shape}, logits range [{logits.min():.4f}, {logits.max():.4f}]" - ) - - return loss, loss_parts - - -def run_one_epoch( - logger: logging.Logger, - args: argparse.Namespace, - model: DeepHealth, - readout: nn.Module | None, - criterion: nn.Module, - train_loader: DataLoader, - optimizer: AdamW, - device: torch.device, - is_train: bool = True, -) -> float: - """ - Run one epoch of training or validation. - - Parameters - ---------- - logger : logging.Logger - args : argparse.Namespace - model : DeepHealth - readout : nn.Module or None - criterion : nn.Module - train_loader : DataLoader - optimizer : AdamW - Unused if is_train=False. - device : torch.device - is_train : bool - If True, perform training updates. Otherwise just evaluate. - - Returns - ------- - avg_loss : float - Average loss over the epoch. - """ - if is_train: - model.train() - else: - model.eval() - - total_loss = 0.0 - n_batches = 0 - n_skipped = 0 - component_sums: Dict[str, float] = {} - - epoch_desc = "train" if is_train else "val" - progress = tqdm( - train_loader, - desc=epoch_desc, - total=len(train_loader), - leave=False, - dynamic_ncols=True, - ) - - for batch_idx, batch in enumerate(progress): - try: - loss, loss_parts = compute_loss( - args=args, - model=model, - readout=readout, - criterion=criterion, - batch=batch, - device=device, - ) - - if is_train: - optimizer.zero_grad(set_to_none=True) - loss.backward() - if args.grad_clip > 0: - clip_grad_norm_(model.parameters(), args.grad_clip) - optimizer.step() - - total_loss += loss.item() - n_batches += 1 - - for name, value in loss_parts.items(): - component_sums[name] = component_sums.get(name, 0.0) + float( - value.detach().item() - ) - - current_loss = loss.item() - avg_loss = total_loss / max(1, n_batches) - postfix = { - "loss": f"{current_loss:.4f}", - "avg": f"{avg_loss:.4f}", - "skipped": n_skipped, - } - for name, sum_value in component_sums.items(): - postfix[name] = f"{sum_value / max(1, n_batches):.4f}" - progress.set_postfix(postfix) - - except RuntimeError as e: - if "Loss is not finite" in str(e): - n_skipped += 1 - logger.warning(f"Batch {batch_idx} skipped: {str(e)[:100]}") - progress.set_postfix( - loss="nan", - avg=f"{total_loss / max(1, n_batches):.4f}" if n_batches > 0 else "inf", - skipped=n_skipped, - ) - continue - else: - raise - - if n_skipped > 0: - logger.info(f"Skipped {n_skipped} batches due to non-finite loss") - - if component_sums and n_batches > 0: - component_summary = ", ".join( - f"{name}={sum_value / n_batches:.4f}" - for name, sum_value in sorted(component_sums.items()) - ) - logger.info(f"Epoch loss breakdown: {component_summary}") - - avg_loss = total_loss / max(1, n_batches) if n_batches > 0 else float("inf") - return avg_loss - - -def evaluate( - logger: logging.Logger, - args: argparse.Namespace, - model: DeepHealth, - readout: nn.Module | None, - criterion: nn.Module, - val_loader: DataLoader, - device: torch.device, -) -> float: - """ - Evaluate model on validation set. - - Returns - ------- - val_loss : float - """ - with torch.no_grad(): - val_loss = run_one_epoch( - logger=logger, - args=args, - model=model, - readout=readout, - criterion=criterion, - train_loader=val_loader, - optimizer=None, - device=device, - is_train=False, - ) - return val_loss - - -def save_checkpoint( - model: DeepHealth, - checkpoint_path: Path, -) -> None: - """Save model state_dict.""" - torch.save(model.state_dict(), checkpoint_path) - - -def save_config( - args: argparse.Namespace, - config_path: Path, - extra: Dict[str, Any] | None = None, -) -> None: - """Save training config as JSON.""" - config_dict: Dict[str, Any] = {} - for key, value in vars(args).items(): - if isinstance(value, tuple): - config_dict[key] = list(value) - elif isinstance(value, list): - config_dict[key] = value - elif isinstance(value, (int, float, str, bool, type(None))): - config_dict[key] = value - else: - config_dict[key] = str(value) - if extra: - config_dict.update(extra) - with open(config_path, "w") as f: - json.dump(config_dict, f, indent=2) - - -def build_run_metadata( - args: argparse.Namespace, - dataset: HealthDataset, - train_subset: Subset, - val_subset: Subset, - test_subset: Subset, - run_name: str, -) -> Dict[str, Any]: - """Collect resolved training facts needed to rebuild the model for evaluation.""" - dataset_class = ( - "AllFutureHealthDataset" - if args.model_target_mode == "all_future" - else "NextStepHealthDataset" - ) - collate_name = ( - "all_future_collate_fn" - if args.model_target_mode == "all_future" - else "next_step_collate_fn" - ) - return { - "run_name": run_name, - "dataset_class": dataset_class, - "collate_fn": collate_name, - "model_class": "DeepHealth", - "model_target_mode": args.model_target_mode, - "dist_mode": args.dist_mode, - "all_future_min_history_events": int(args.all_future_min_history_events), - "all_future_min_future_events": int(args.all_future_min_future_events), - "extra_info_types_file": ( - Path(args.extra_info_types_file).name - if args.extra_info_types_file is not None - else None - ), - "extra_info_types": dataset.extra_info_types, - "dataset_metadata": { - "vocab_size": int(dataset.vocab_size), - "n_types": int(dataset.n_types), - "n_cont_types": int(dataset.n_cont_types), - "n_categories": int(dataset.n_categories), - "cont_type_ids": [int(x) for x in dataset.cont_type_ids], - "extra_info_types": [int(x) for x in dataset.extra_info_types], - }, - "split_sizes": { - "train": int(len(train_subset)), - "val": int(len(val_subset)), - "test": int(len(test_subset)), - }, - "resolved_readout_name": args.readout_name, - "resolved_loss_name": args.loss_name, - } - - -def normalize_training_config(args: argparse.Namespace) -> None: - """Fill in and validate training options that depend on other flags.""" - if args.model_target_mode not in {"next_token", "all_future"}: - raise ValueError(f"Unknown model_target_mode: {args.model_target_mode}") - if args.dist_mode not in {"exponential", "weibull", "mixed"}: - raise ValueError(f"Unknown dist_mode: {args.dist_mode}") - if args.model_target_mode == "next_token" and args.dist_mode != "exponential": - raise ValueError( - "next_token training currently supports dist_mode='exponential' only. " - "Use model_target_mode='all_future' for weibull or mixed distributions." - ) - if args.all_future_min_history_events < 1: - raise ValueError("all_future_min_history_events must be >= 1") - if args.all_future_min_future_events < 1: - raise ValueError("all_future_min_future_events must be >= 1") - if args.model_target_mode == "all_future": - args.target_mode = "all_future" - - # gap_5y is always enabled, so preserve NO_EVENT target behavior. - args.ignore_no_event_in_delphi2m = False - if args.model_target_mode == "next_token" and args.target_mode == "uts": - args.include_no_event_in_uts_target = True - - -def normalize_loss_and_distribution_config(args: argparse.Namespace) -> None: - """Validate and resolve loss/distribution options after auto-selection.""" - next_token_losses = {"delphi2m", "uts"} - all_future_losses = {"exponential", "weibull", "mixed"} - - if args.model_target_mode == "all_future": - if args.loss_name not in all_future_losses: - raise ValueError( - "all_future training requires loss_name to be one of " - "exponential, weibull, or mixed." - ) - if args.loss_name != args.dist_mode: - raise ValueError( - "all_future loss_name must match dist_mode so risk scoring and " - f"training distribution stay aligned. Got loss_name={args.loss_name!r}, " - f"dist_mode={args.dist_mode!r}." - ) - return - - if args.target_mode not in {"delphi2m", "uts"}: - raise ValueError(f"Unknown target_mode: {args.target_mode}") - - if args.loss_name not in next_token_losses: - raise ValueError( - "Unknown loss_name. Supported values: delphi2m, uts." - ) - - if args.loss_name == "delphi2m" and args.target_mode != "delphi2m": - raise ValueError( - "loss_name=delphi2m requires target_mode=delphi2m." - ) - - if args.loss_name == "uts" and args.target_mode != "uts": - raise ValueError( - "loss_name=uts requires target_mode=uts." - ) - - -# --------------------------------------------------------------------------- -# Main Training Loop -# --------------------------------------------------------------------------- - -def main(): - """Main training function.""" - parser = argparse.ArgumentParser(description="DeepHealth Training") - - # ---- Data & Output ---- - parser.add_argument("--data_prefix", type=str, default="ukb", - help="Prefix for data files") - parser.add_argument("--labels_file", type=str, default="labels.csv", - help="Path to labels file") - parser.add_argument("--seed", type=int, default=42, - help="Random seed") - - # ---- Dataset ---- - parser.add_argument("--no_event_interval_years", type=float, default=5.0, - help="Interval in years for no-event insertion") - parser.add_argument("--include_no_event_in_uts_target", action="store_true", - help="Include NO_EVENT in UTS target multi-hot") - - # ---- Split ---- - parser.add_argument("--train_ratio", type=float, default=0.7, - help="Training set ratio") - parser.add_argument("--val_ratio", type=float, default=0.15, - help="Validation set ratio") - parser.add_argument("--test_ratio", type=float, default=0.15, - help="Test set ratio") - - # ---- Model ---- - parser.add_argument("--n_embd", type=int, default=120, - help="Embedding dimension") - parser.add_argument("--n_head", type=int, default=10, - help="Number of attention heads") - parser.add_argument("--n_hist_layer", type=int, default=12, - help="Number of history encoder layers") - parser.add_argument("--n_tab_layer", type=int, default=4, - help="Number of self-attention layers for other-token encoder") - parser.add_argument("--n_bins", type=int, default=16, - help="Number of bins for continuous other-token values") - parser.add_argument("--time_mode", type=str, default="relative", - choices=["relative", "absolute"], - help="Time encoding mode for disease history") - parser.add_argument("--model_target_mode", type=str, default="next_token", - choices=["next_token", "all_future"], - help="Model forward/training mode") - parser.add_argument("--dist_mode", type=str, default="exponential", - choices=["exponential", "weibull", "mixed"], - help="Event-time distribution. next_token requires exponential; all_future supports exponential, weibull, and mixed") - parser.add_argument("--dropout", type=float, default=0.0, - help="Dropout rate") - parser.add_argument("--extra_info_types_file", type=str, default=None, - help="Optional file containing other-information type ids to include") - - # ---- Training Protocol ---- - parser.add_argument("--target_mode", type=str, default="uts", - choices=["delphi2m", "uts"], - help="Next-token supervision mode; ignored for all_future model_target_mode") - parser.add_argument("--readout_name", type=str, default=None, - help="Readout name (auto-selected if None)") - parser.add_argument("--readout_reduce", type=str, default="mean", - choices=["mean", "sum"], - help="Readout reduction for SameTimeGroupEndReadout") - parser.add_argument("--all_future_min_history_events", type=int, default=1, - help="Minimum historical events before an all-future query") - parser.add_argument("--all_future_min_future_events", type=int, default=1, - help="Minimum future events after an all-future query") - - # ---- Loss ---- - parser.add_argument("--loss_name", type=str, default=None, - help="Loss name (auto-selected if None): delphi2m, uts, exponential, weibull, mixed") - parser.add_argument("--t_min", type=float, default=0.0027378507871321013, - help="Minimum time for loss (1/365.25)") - parser.add_argument("--max_exp_input", type=float, default=60.0, - help="Max exponent input for loss") - parser.add_argument("--ce_weight", type=float, default=1.0, - help="Cross-entropy weight in delphi2m loss") - parser.add_argument("--time_weight", type=float, default=1.0, - help="Time loss weight in delphi2m loss") - parser.add_argument("--ignore_no_event_in_delphi2m", action="store_true", - help="Ignore NO_EVENT in delphi2m loss") - - # ---- Optimization ---- - parser.add_argument("--batch_size", type=int, default=128, - help="Batch size") - parser.add_argument("--base_lr", type=float, default=3e-4, - help="Base learning rate") - parser.add_argument("--weight_decay", type=float, default=0.1, - help="Weight decay (L2 regularization)") - parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99), - help="AdamW betas") - parser.add_argument("--grad_clip", type=float, default=1.0, - help="Gradient clipping norm") - parser.add_argument("--max_epochs", type=int, default=200, - help="Maximum number of epochs") - parser.add_argument("--warmup_epochs", type=int, default=10, - help="Number of warmup epochs") - parser.add_argument("--patience", type=int, default=15, - help="Early stopping patience") - parser.add_argument("--min_lr_ratio", type=float, default=0.1, - help="Minimum LR as ratio of adaptive_lr") - parser.add_argument("--num_workers", type=int, default=4, - help="Number of DataLoader workers") - parser.add_argument("--device", type=str, default="cuda", - help="Device to use for training: cpu, cuda, or cuda:") - - args = parser.parse_args() - args.extra_info_types = ( - load_extra_info_types_file(args.extra_info_types_file) - if args.extra_info_types_file is not None - else None - ) - - # ---- Setup ---- - set_seed(args.seed) - device = resolve_device(args.device) - configure_torch_for_training(device) - - normalize_training_config(args) - # Auto-select readout if not specified. - if args.model_target_mode == "all_future": - args.readout_name = "none" - elif args.readout_name is None: - args.readout_name = ( - "token" if args.target_mode == "delphi2m" - else "same_time_group_end" - ) - - # Auto-select loss if not specified. - if args.loss_name is None: - if args.model_target_mode == "all_future": - args.loss_name = args.dist_mode - else: - args.loss_name = ( - "delphi2m" if args.target_mode == "delphi2m" - else "uts" - ) - - normalize_loss_and_distribution_config(args) - - runs_root = Path("runs") - while True: - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - run_name = ( - f"{args.time_mode}_{args.dist_mode}_{args.model_target_mode}_{args.loss_name}_" - f"other_tokens_gap_5y_{timestamp}" - ) - run_dir = runs_root / run_name - if not run_dir.exists(): - break - time.sleep(1) - - run_dir.mkdir(parents=True, exist_ok=False) - - logger = setup_logging(run_dir) - logger.info(f"Starting training run: {run_name}") - logger.info(f"Device: {device}") - logger.info(f"extra_info_types: {args.extra_info_types or 'all'}") - logger.info( - f"Resolved no_event config: ignore_no_event_in_delphi2m={args.ignore_no_event_in_delphi2m}, " - f"include_no_event_in_uts_target={args.include_no_event_in_uts_target}" - ) - - # ---- Validate Arguments ---- - total_ratio = args.train_ratio + args.val_ratio + args.test_ratio - if not np.isclose(total_ratio, 1.0, atol=1e-6): - raise ValueError( - f"train_ratio + val_ratio + test_ratio must equal 1.0, got {total_ratio}" - ) - - logger.info(f"Auto-selected readout: {args.readout_name}") - logger.info(f"Auto-selected loss: {args.loss_name}") - - # ---- Load Dataset ---- - logger.info("Loading dataset...") - if args.model_target_mode == "all_future": - dataset = AllFutureHealthDataset( - data_prefix=args.data_prefix, - labels_file=args.labels_file, - split="train", - no_event_interval_years=args.no_event_interval_years, - include_no_event_in_uts_target=args.include_no_event_in_uts_target, - min_history_events=args.all_future_min_history_events, - min_future_events=args.all_future_min_future_events, - extra_info_types=args.extra_info_types, - ) - val_dataset = AllFutureHealthDataset( - data_prefix=args.data_prefix, - labels_file=args.labels_file, - split="valid", - no_event_interval_years=args.no_event_interval_years, - include_no_event_in_uts_target=args.include_no_event_in_uts_target, - min_history_events=args.all_future_min_history_events, - min_future_events=args.all_future_min_future_events, - extra_info_types=args.extra_info_types, - ) - test_dataset = AllFutureHealthDataset( - data_prefix=args.data_prefix, - labels_file=args.labels_file, - split="test", - no_event_interval_years=args.no_event_interval_years, - include_no_event_in_uts_target=args.include_no_event_in_uts_target, - min_history_events=args.all_future_min_history_events, - min_future_events=args.all_future_min_future_events, - extra_info_types=args.extra_info_types, - ) - train_subset, val_subset, test_subset = split_all_future_datasets( - train_dataset=dataset, - val_dataset=val_dataset, - test_dataset=test_dataset, - train_ratio=args.train_ratio, - val_ratio=args.val_ratio, - test_ratio=args.test_ratio, - seed=args.seed, - ) - active_collate_fn = all_future_collate_fn - else: - dataset = HealthDataset( - data_prefix=args.data_prefix, - labels_file=args.labels_file, - no_event_interval_years=args.no_event_interval_years, - include_no_event_in_uts_target=args.include_no_event_in_uts_target, - extra_info_types=args.extra_info_types, - ) - train_subset, val_subset, test_subset = split_dataset( - dataset=dataset, - train_ratio=args.train_ratio, - val_ratio=args.val_ratio, - test_ratio=args.test_ratio, - seed=args.seed, - ) - active_collate_fn = collate_fn - - logger.info( - f"Dataset loaded: {len(dataset)} base samples, vocab_size={dataset.vocab_size}") - - logger.info( - f"Dataset split: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}" - ) - - # ---- Build DataLoaders ---- - train_loader = DataLoader( - train_subset, - batch_size=args.batch_size, - sampler=RandomSampler( - train_subset, generator=torch.Generator().manual_seed(args.seed)), - collate_fn=active_collate_fn, - num_workers=args.num_workers, - pin_memory=device.type == "cuda", - persistent_workers=args.num_workers > 0, - prefetch_factor=2 if args.num_workers > 0 else None, - ) - val_loader = DataLoader( - val_subset, - batch_size=args.batch_size, - shuffle=False, - collate_fn=active_collate_fn, - num_workers=args.num_workers, - pin_memory=device.type == "cuda", - persistent_workers=args.num_workers > 0, - prefetch_factor=2 if args.num_workers > 0 else None, - ) - test_loader = DataLoader( - test_subset, - batch_size=args.batch_size, - shuffle=False, - collate_fn=active_collate_fn, - num_workers=args.num_workers, - pin_memory=device.type == "cuda", - persistent_workers=args.num_workers > 0, - prefetch_factor=2 if args.num_workers > 0 else None, - ) - - # ---- Build Model ---- - logger.info("Building model...") - model = build_model(args, dataset) - model.to(device) - n_params = sum(p.numel() for p in model.parameters()) - logger.info(f"Model built: {n_params:,} parameters") - - # ---- Build Optimizer ---- - optimizer = build_optimizer(args, model) - logger.info( - f"Optimizer: AdamW, base_lr={args.base_lr}, weight_decay={args.weight_decay}") - - # ---- Compute Adaptive LR ---- - adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128) - logger.info( - f"Adaptive LR: {adaptive_lr:.6f} (base_lr * sqrt(batch_size/128))") - - # ---- Build Readout ---- - if args.model_target_mode == "all_future": - readout = None - elif args.readout_name == "token": - readout = build_readout("token") - elif args.readout_name == "same_time_group_end": - readout = build_readout("same_time_group_end", - reduce=args.readout_reduce) - elif args.readout_name == "last_valid": - readout = build_readout("last_valid") - else: - raise ValueError(f"Unknown readout: {args.readout_name}") - logger.info(f"Readout: {args.readout_name}") - - # ---- Build Loss ---- - if args.model_target_mode == "all_future": - ignored_idx = {PAD_IDX, CHECKUP_IDX} - if args.loss_name == "exponential": - criterion = build_loss("exponential", ignored_idx=ignored_idx) - elif args.loss_name == "weibull": - criterion = build_loss("weibull", ignored_idx=ignored_idx) - elif args.loss_name == "mixed": - criterion = build_loss( - "mixed", - death_idx=dataset.vocab_size - 1, - ignored_idx=ignored_idx, - ) - else: - raise ValueError(f"Unknown all_future loss: {args.loss_name}") - logger.info( - f"Loss: {args.loss_name}, dist_mode={args.dist_mode}, ignored_idx={ignored_idx}") - elif args.loss_name == "delphi2m": - ignored_tokens = {PAD_IDX, CHECKUP_IDX} - if args.ignore_no_event_in_delphi2m: - ignored_tokens.add(NO_EVENT_IDX) - criterion = build_loss( - "delphi2m", - ignored_tokens=ignored_tokens, - t_min=args.t_min, - max_exp_input=args.max_exp_input, - ce_weight=args.ce_weight, - time_weight=args.time_weight, - ) - logger.info(f"Loss: delphi2m, ignored_tokens={ignored_tokens}") - elif args.loss_name == "uts": - ignored_idx = {PAD_IDX, CHECKUP_IDX} - criterion = build_loss( - "uts", - ignored_idx=ignored_idx, - t_min=args.t_min, - max_exp_input=args.max_exp_input, - ) - logger.info(f"Loss: uts, ignored_idx={ignored_idx}") - else: - raise ValueError(f"Unknown loss: {args.loss_name}") - - # ---- Save Config ---- - save_config( - args, - run_dir / "train_config.json", - extra=build_run_metadata( - args=args, - dataset=dataset, - train_subset=train_subset, - val_subset=val_subset, - test_subset=test_subset, - run_name=run_name, - ), - ) - logger.info(f"Config saved to {run_dir / 'train_config.json'}") - - # ---- Training Loop ---- - logger.info("Starting training...") - best_val_loss = float("inf") - patience_counter = 0 - metrics = [] - - best_model_path = run_dir / "best_model.pt" - history_path = run_dir / "history.json" - - start_time = time.time() - - for epoch in range(args.max_epochs): - lr = get_lr(epoch, args, adaptive_lr, args.max_epochs) - set_optimizer_lr(optimizer, lr) - - # Train - train_loss = run_one_epoch( - logger=logger, - args=args, - model=model, - readout=readout, - criterion=criterion, - train_loader=train_loader, - optimizer=optimizer, - device=device, - is_train=True, - ) - - # Validate - val_loss = evaluate( - logger=logger, - args=args, - model=model, - readout=readout, - criterion=criterion, - val_loader=val_loader, - device=device, - ) - - # Early stopping - is_best = False - if val_loss < best_val_loss: - best_val_loss = val_loss - patience_counter = 0 - is_best = True - save_checkpoint(model, best_model_path) - else: - patience_counter += 1 - - elapsed = time.time() - start_time - logger.info( - f"Epoch {epoch+1}/{args.max_epochs} | lr={lr:.6f} | " - f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | " - f"best_val_loss={best_val_loss:.6f} | patience={patience_counter}/{args.patience} | " - f"elapsed={elapsed:.1f}s" - ) - - metrics.append({ - "epoch": epoch + 1, - "lr": lr, - "train_loss": train_loss, - "val_loss": val_loss, - "best_val_loss": best_val_loss, - "is_best": int(is_best), - }) - - if patience_counter >= args.patience: - logger.info(f"Early stopping triggered at epoch {epoch+1}") - break - - # ---- Save Training History ---- - with open(history_path, "w") as f: - json.dump(metrics, f, indent=2) - logger.info(f"History saved to {history_path}") - - # ---- Test on Best Model ---- - logger.info("Evaluating best model on test set...") - best_state_dict = torch.load(best_model_path, map_location=device) - model.load_state_dict(best_state_dict) - - test_loss = evaluate( - logger=logger, - args=args, - model=model, - readout=readout, - criterion=criterion, - val_loader=test_loader, - device=device, - ) - logger.info(f"Test loss: {test_loss:.6f}") - - total_time = time.time() - start_time - logger.info(f"Training completed in {total_time:.1f}s") - logger.info(f"Best checkpoint: {best_model_path}") - - -if __name__ == "__main__": - main() diff --git a/train_all_future.py b/train_all_future.py new file mode 100644 index 0000000..e8d96ea --- /dev/null +++ b/train_all_future.py @@ -0,0 +1,448 @@ +""" +Train DeepHealth with query-conditioned all-future supervision. + +Training samples are patient-level. For each patient and each __getitem__ call, +AllFutureHealthDataset randomly samples a query time t_query, uses events at or +before t_query as history, and uses events after t_query as the future target set. + +Validation/test samples are deterministic query points built from future event +times, then split by patient. +""" +from __future__ import annotations + +import argparse +import json +import logging +import math +import time +from datetime import datetime +from pathlib import Path +from typing import Any, Dict + +import numpy as np +import torch +from torch.nn.utils import clip_grad_norm_ +from torch.optim import AdamW +from torch.utils.data import DataLoader, RandomSampler +from tqdm.auto import tqdm + +from dataset import AllFutureHealthDataset, all_future_collate_fn +from losses import build_loss +from models import DeepHealth +from targets import CHECKUP_IDX, PAD_IDX +from train_util import ( + configure_torch_for_training, + load_extra_info_types_file, + resolve_device, + save_checkpoint, + save_config, + set_optimizer_lr, + set_seed, + setup_logging, + split_all_future_datasets, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Train DeepHealth with all-future supervision") + + parser.add_argument("--data_prefix", type=str, default="ukb") + parser.add_argument("--labels_file", type=str, default="labels.csv") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--extra_info_types_file", type=str, default=None) + + parser.add_argument("--train_ratio", type=float, default=0.7) + parser.add_argument("--val_ratio", type=float, default=0.15) + parser.add_argument("--test_ratio", type=float, default=0.15) + parser.add_argument("--min_history_events", type=int, default=1) + parser.add_argument("--min_future_events", type=int, default=1) + parser.add_argument("--validation_query_seed", type=int, default=None) + + parser.add_argument("--n_embd", type=int, default=120) + parser.add_argument("--n_head", type=int, default=10) + parser.add_argument("--n_hist_layer", type=int, default=12) + parser.add_argument("--n_tab_layer", type=int, default=4) + parser.add_argument("--n_bins", type=int, default=16) + parser.add_argument("--time_mode", type=str, default="relative", + choices=["relative", "absolute"]) + parser.add_argument("--dist_mode", type=str, default="exponential", + choices=["exponential", "weibull", "mixed"]) + parser.add_argument("--dropout", type=float, default=0.0) + + parser.add_argument("--batch_size", type=int, default=128) + parser.add_argument("--base_lr", type=float, default=3e-4) + parser.add_argument("--weight_decay", type=float, default=0.1) + parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99)) + parser.add_argument("--grad_clip", type=float, default=1.0) + parser.add_argument("--max_epochs", type=int, default=200) + parser.add_argument("--warmup_epochs", type=int, default=10) + parser.add_argument("--patience", type=int, default=15) + parser.add_argument("--min_lr_ratio", type=float, default=0.1) + parser.add_argument("--num_workers", type=int, default=4) + parser.add_argument("--device", type=str, default="cuda") + + args = parser.parse_args() + if args.min_history_events < 1: + raise ValueError("min_history_events must be >= 1") + if args.min_future_events < 1: + raise ValueError("min_future_events must be >= 1") + if not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0): + raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0") + if args.validation_query_seed is None: + args.validation_query_seed = int(args.seed) + args.extra_info_types = ( + load_extra_info_types_file(args.extra_info_types_file) + if args.extra_info_types_file is not None + else None + ) + return args + + +def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float: + if epoch < args.warmup_epochs: + return adaptive_lr * (epoch + 1) / args.warmup_epochs + progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs) + cosine = 0.5 * (1 + math.cos(math.pi * progress)) + return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio)) + + +def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]: + non_blocking = device.type == "cuda" + return { + key: value.to(device, non_blocking=non_blocking) + if isinstance(value, torch.Tensor) + else value + for key, value in batch.items() + } + + +def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> DeepHealth: + return DeepHealth( + vocab_size=dataset.vocab_size, + n_embd=args.n_embd, + n_head=args.n_head, + n_hist_layer=args.n_hist_layer, + n_tab_layer=args.n_tab_layer, + n_types=dataset.n_types, + n_cont_types=dataset.n_cont_types, + n_categories=dataset.n_categories, + cont_type_ids=dataset.cont_type_ids, + n_bins=args.n_bins, + target_mode="all_future", + time_mode=args.time_mode, + dist_mode=args.dist_mode, + dropout=args.dropout, + ) + + +def build_criterion(args: argparse.Namespace, dataset: AllFutureHealthDataset): + ignored_idx = {PAD_IDX, CHECKUP_IDX} + if args.dist_mode == "exponential": + return build_loss("exponential", ignored_idx=ignored_idx) + if args.dist_mode == "weibull": + return build_loss("weibull", ignored_idx=ignored_idx) + if args.dist_mode == "mixed": + return build_loss( + "mixed", + death_idx=dataset.vocab_size - 1, + ignored_idx=ignored_idx, + ) + raise ValueError(f"Unknown dist_mode: {args.dist_mode}") + + +def compute_all_future_loss( + args: argparse.Namespace, + model: DeepHealth, + criterion, + batch: Dict[str, torch.Tensor], + device: torch.device, +) -> torch.Tensor: + batch = move_batch_to_device(batch, device) + + hidden = model( + event_seq=batch["event_seq"], + time_seq=batch["time_seq"], + sex=batch["sex"], + padding_mask=batch["padding_mask"], + t_query=batch["t_query"], + other_type=batch["other_type"], + other_value=batch["other_value"], + other_value_kind=batch["other_value_kind"], + other_time=batch["other_time"], + target_mode="all_future", + ) + logits = model.calc_risk(hidden) + + if args.dist_mode == "exponential": + loss = criterion( + logits=logits, + targets=batch["future_targets"], + exposure=batch["exposure"], + ) + elif args.dist_mode == "weibull": + loss = criterion( + logits=logits, + weibull_rho=model.calc_weibull_rho(hidden), + targets=batch["future_targets"], + dt=batch["future_dt"], + exposure=batch["exposure"], + ) + else: + loss = criterion( + logits=logits, + death_rho=model.calc_death_rho(hidden), + targets=batch["future_targets"], + dt=batch["future_dt"], + exposure=batch["exposure"], + ) + + if not torch.isfinite(loss): + raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}") + return loss + + +def run_epoch( + logger: logging.Logger, + args: argparse.Namespace, + model: DeepHealth, + criterion, + loader: DataLoader, + optimizer: AdamW | None, + device: torch.device, + is_train: bool, +) -> float: + model.train(is_train) + total = 0.0 + n_batches = 0 + skipped = 0 + desc = "train" if is_train else "val" + + progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True) + for batch_idx, batch in enumerate(progress): + try: + loss = compute_all_future_loss(args, model, criterion, batch, device) + if is_train: + if optimizer is None: + raise ValueError("optimizer is required for training") + optimizer.zero_grad(set_to_none=True) + loss.backward() + if args.grad_clip > 0: + clip_grad_norm_(model.parameters(), args.grad_clip) + optimizer.step() + + total += float(loss.detach().cpu()) + n_batches += 1 + avg = total / max(1, n_batches) + progress.set_postfix(loss=f"{float(loss.detach().cpu()):.4f}", avg=f"{avg:.4f}", skipped=skipped) + except RuntimeError as exc: + if "Loss is not finite" not in str(exc): + raise + skipped += 1 + logger.warning(f"Batch {batch_idx} skipped: {str(exc)[:120]}") + + if skipped: + logger.info(f"Skipped {skipped} batches due to non-finite loss") + return total / max(1, n_batches) if n_batches else float("inf") + + +def build_metadata( + args: argparse.Namespace, + dataset: AllFutureHealthDataset, + run_name: str, + train_subset, + val_subset, + test_subset, +) -> Dict[str, Any]: + return { + "run_name": run_name, + "dataset_class": "AllFutureHealthDataset", + "collate_fn": "all_future_collate_fn", + "model_class": "DeepHealth", + "model_target_mode": "all_future", + "target_mode": "all_future", + "dist_mode": args.dist_mode, + "all_future_min_history_events": int(args.min_history_events), + "all_future_min_future_events": int(args.min_future_events), + "all_future_validation_query_seed": int(args.validation_query_seed), + "extra_info_types_file": ( + Path(args.extra_info_types_file).name + if args.extra_info_types_file is not None + else None + ), + "extra_info_types": [int(x) for x in dataset.extra_info_types], + "dataset_metadata": { + "vocab_size": int(dataset.vocab_size), + "n_types": int(dataset.n_types), + "n_cont_types": int(dataset.n_cont_types), + "n_categories": int(dataset.n_categories), + "cont_type_ids": [int(x) for x in dataset.cont_type_ids], + "extra_info_types": [int(x) for x in dataset.extra_info_types], + }, + "split_sizes": { + "train": int(len(train_subset)), + "val": int(len(val_subset)), + "test": int(len(test_subset)), + }, + "resolved_readout_name": "none", + "resolved_loss_name": args.dist_mode, + } + + +def main() -> None: + args = parse_args() + set_seed(args.seed) + device = resolve_device(args.device) + configure_torch_for_training(device) + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + run_name = f"{args.time_mode}_{args.dist_mode}_all_future_pure_disease_{timestamp}" + run_dir = Path("runs") / run_name + run_dir.mkdir(parents=True, exist_ok=False) + logger = setup_logging(run_dir) + + logger.info(f"Starting all-future training run: {run_name}") + logger.info(f"Device: {device}") + logger.info(f"extra_info_types: {args.extra_info_types or 'all'}") + + logger.info("Loading all-future datasets...") + train_dataset = AllFutureHealthDataset( + data_prefix=args.data_prefix, + labels_file=args.labels_file, + split="train", + min_history_events=args.min_history_events, + min_future_events=args.min_future_events, + validation_query_seed=args.validation_query_seed, + extra_info_types=args.extra_info_types, + ) + val_dataset = AllFutureHealthDataset( + data_prefix=args.data_prefix, + labels_file=args.labels_file, + split="valid", + min_history_events=args.min_history_events, + min_future_events=args.min_future_events, + validation_query_seed=args.validation_query_seed, + extra_info_types=args.extra_info_types, + ) + test_dataset = AllFutureHealthDataset( + data_prefix=args.data_prefix, + labels_file=args.labels_file, + split="test", + min_history_events=args.min_history_events, + min_future_events=args.min_future_events, + validation_query_seed=args.validation_query_seed, + extra_info_types=args.extra_info_types, + ) + train_subset, val_subset, test_subset = split_all_future_datasets( + train_dataset=train_dataset, + val_dataset=val_dataset, + test_dataset=test_dataset, + train_ratio=args.train_ratio, + val_ratio=args.val_ratio, + test_ratio=args.test_ratio, + seed=args.seed, + ) + logger.info( + f"Patients/queries: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}" + ) + + train_loader = DataLoader( + train_subset, + batch_size=args.batch_size, + sampler=RandomSampler(train_subset, generator=torch.Generator().manual_seed(args.seed)), + collate_fn=all_future_collate_fn, + num_workers=args.num_workers, + pin_memory=device.type == "cuda", + persistent_workers=args.num_workers > 0, + prefetch_factor=2 if args.num_workers > 0 else None, + ) + val_loader = DataLoader( + val_subset, + batch_size=args.batch_size, + shuffle=False, + collate_fn=all_future_collate_fn, + num_workers=args.num_workers, + pin_memory=device.type == "cuda", + persistent_workers=args.num_workers > 0, + prefetch_factor=2 if args.num_workers > 0 else None, + ) + test_loader = DataLoader( + test_subset, + batch_size=args.batch_size, + shuffle=False, + collate_fn=all_future_collate_fn, + num_workers=args.num_workers, + pin_memory=device.type == "cuda", + persistent_workers=args.num_workers > 0, + prefetch_factor=2 if args.num_workers > 0 else None, + ) + + model = build_model(args, train_dataset).to(device) + optimizer = AdamW( + model.parameters(), + lr=args.base_lr, + betas=tuple(args.betas), + weight_decay=args.weight_decay, + ) + criterion = build_criterion(args, train_dataset) + adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128) + + save_config( + args, + run_dir / "train_config.json", + extra=build_metadata(args, train_dataset, run_name, train_subset, val_subset, test_subset), + ) + + best_val = float("inf") + patience = 0 + history = [] + best_model_path = run_dir / "best_model.pt" + start = time.time() + + for epoch in range(args.max_epochs): + lr = get_lr(epoch, args, adaptive_lr) + set_optimizer_lr(optimizer, lr) + + train_loss = run_epoch(logger, args, model, criterion, train_loader, optimizer, device, True) + with torch.no_grad(): + val_loss = run_epoch(logger, args, model, criterion, val_loader, None, device, False) + + is_best = val_loss < best_val + if is_best: + best_val = val_loss + patience = 0 + save_checkpoint(model, best_model_path) + else: + patience += 1 + + logger.info( + f"Epoch {epoch + 1}/{args.max_epochs} | lr={lr:.6f} | " + f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | " + f"best_val_loss={best_val:.6f} | patience={patience}/{args.patience} | " + f"elapsed={time.time() - start:.1f}s" + ) + history.append({ + "epoch": epoch + 1, + "lr": lr, + "train_loss": train_loss, + "val_loss": val_loss, + "best_val_loss": best_val, + "is_best": int(is_best), + }) + if patience >= args.patience: + logger.info(f"Early stopping triggered at epoch {epoch + 1}") + break + + with (run_dir / "history.json").open("w", encoding="utf-8") as f: + json.dump(history, f, indent=2) + + logger.info("Evaluating best model on all-future test queries...") + model.load_state_dict(torch.load(best_model_path, map_location=device)) + with torch.no_grad(): + test_loss = run_epoch(logger, args, model, criterion, test_loader, None, device, False) + logger.info(f"Test loss: {test_loss:.6f}") + logger.info(f"Best checkpoint: {best_model_path}") + + +if __name__ == "__main__": + main() diff --git a/train_next_step.py b/train_next_step.py new file mode 100644 index 0000000..02e044c --- /dev/null +++ b/train_next_step.py @@ -0,0 +1,454 @@ +""" +Train DeepHealth with next-token / next-time-point supervision. + +The dataset remains the current next-step construction: pure disease events plus +optional gap imputation are shifted into autoregressive inputs and +targets. UTS training reads out only same-time group ends. +""" +from __future__ import annotations + +import argparse +import json +import logging +import math +import time +from datetime import datetime +from pathlib import Path +from typing import Any, Dict + +import numpy as np +import torch +from torch.nn.utils import clip_grad_norm_ +from torch.optim import AdamW +from torch.utils.data import DataLoader, RandomSampler +from tqdm.auto import tqdm + +from dataset import HealthDataset, collate_fn +from losses import build_loss +from models import DeepHealth +from readouts import build_readout +from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX +from train_util import ( + configure_torch_for_training, + load_extra_info_types_file, + resolve_device, + save_checkpoint, + save_config, + set_optimizer_lr, + set_seed, + setup_logging, + split_dataset, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Train DeepHealth with next-token/point supervision") + + parser.add_argument("--data_prefix", type=str, default="ukb") + parser.add_argument("--labels_file", type=str, default="labels.csv") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--extra_info_types_file", type=str, default=None) + parser.add_argument("--no_event_interval_years", type=float, default=5.0) + parser.add_argument("--include_no_event_in_uts_target", action="store_true") + + parser.add_argument("--train_ratio", type=float, default=0.7) + parser.add_argument("--val_ratio", type=float, default=0.15) + parser.add_argument("--test_ratio", type=float, default=0.15) + + parser.add_argument("--n_embd", type=int, default=120) + parser.add_argument("--n_head", type=int, default=10) + parser.add_argument("--n_hist_layer", type=int, default=12) + parser.add_argument("--n_tab_layer", type=int, default=4) + parser.add_argument("--n_bins", type=int, default=16) + parser.add_argument("--time_mode", type=str, default="relative", + choices=["relative", "absolute"]) + parser.add_argument("--dropout", type=float, default=0.0) + + parser.add_argument("--target_mode", type=str, default="uts", + choices=["delphi2m", "uts"]) + parser.add_argument("--readout_name", type=str, default=None, + choices=["token", "same_time_group_end", "last_valid"]) + parser.add_argument("--readout_reduce", type=str, default="mean", + choices=["mean", "sum"]) + parser.add_argument("--t_min", type=float, default=0.0027378507871321013) + parser.add_argument("--max_exp_input", type=float, default=60.0) + parser.add_argument("--ce_weight", type=float, default=1.0) + parser.add_argument("--time_weight", type=float, default=1.0) + parser.add_argument("--ignore_no_event_in_delphi2m", action="store_true") + + parser.add_argument("--batch_size", type=int, default=128) + parser.add_argument("--base_lr", type=float, default=3e-4) + parser.add_argument("--weight_decay", type=float, default=0.1) + parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99)) + parser.add_argument("--grad_clip", type=float, default=1.0) + parser.add_argument("--max_epochs", type=int, default=200) + parser.add_argument("--warmup_epochs", type=int, default=10) + parser.add_argument("--patience", type=int, default=15) + parser.add_argument("--min_lr_ratio", type=float, default=0.1) + parser.add_argument("--num_workers", type=int, default=4) + parser.add_argument("--device", type=str, default="cuda") + + args = parser.parse_args() + if not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0): + raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0") + if args.target_mode == "uts": + args.readout_name = args.readout_name or "same_time_group_end" + args.include_no_event_in_uts_target = True + else: + args.readout_name = args.readout_name or "token" + args.extra_info_types = ( + load_extra_info_types_file(args.extra_info_types_file) + if args.extra_info_types_file is not None + else None + ) + return args + + +def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float: + if epoch < args.warmup_epochs: + return adaptive_lr * (epoch + 1) / args.warmup_epochs + progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs) + cosine = 0.5 * (1 + math.cos(math.pi * progress)) + return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio)) + + +def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]: + non_blocking = device.type == "cuda" + return { + key: value.to(device, non_blocking=non_blocking) + if isinstance(value, torch.Tensor) + else value + for key, value in batch.items() + } + + +def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth: + return DeepHealth( + vocab_size=dataset.vocab_size, + n_embd=args.n_embd, + n_head=args.n_head, + n_hist_layer=args.n_hist_layer, + n_tab_layer=args.n_tab_layer, + n_types=dataset.n_types, + n_cont_types=dataset.n_cont_types, + n_categories=dataset.n_categories, + cont_type_ids=dataset.cont_type_ids, + n_bins=args.n_bins, + target_mode="next_token", + time_mode=args.time_mode, + dist_mode="exponential", + dropout=args.dropout, + ) + + +def build_next_step_readout(args: argparse.Namespace): + if args.readout_name == "same_time_group_end": + return build_readout("same_time_group_end", reduce=args.readout_reduce) + return build_readout(args.readout_name) + + +def build_next_step_loss(args: argparse.Namespace): + if args.target_mode == "delphi2m": + ignored_tokens = {PAD_IDX, CHECKUP_IDX} + if args.ignore_no_event_in_delphi2m: + ignored_tokens.add(NO_EVENT_IDX) + return build_loss( + "delphi2m", + ignored_tokens=ignored_tokens, + t_min=args.t_min, + max_exp_input=args.max_exp_input, + ce_weight=args.ce_weight, + time_weight=args.time_weight, + ) + return build_loss( + "uts", + ignored_idx={PAD_IDX, CHECKUP_IDX}, + t_min=args.t_min, + max_exp_input=args.max_exp_input, + ) + + +def compute_next_step_loss( + args: argparse.Namespace, + model: DeepHealth, + readout, + criterion, + batch: Dict[str, torch.Tensor], + device: torch.device, +) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]: + batch = move_batch_to_device(batch, device) + hidden = model( + event_seq=batch["event_seq"], + time_seq=batch["time_seq"], + sex=batch["sex"], + padding_mask=batch["padding_mask"], + other_type=batch["other_type"], + other_value=batch["other_value"], + other_value_kind=batch["other_value_kind"], + other_time=batch["other_time"], + target_mode="next_token", + ) + readout_out = readout( + hidden=hidden, + time_seq=batch["time_seq"], + padding_mask=batch["padding_mask"], + readout_mask=batch["readout_mask"] + if args.readout_name == "same_time_group_end" + else None, + ) + logits = model.calc_risk(readout_out.hidden) + + if args.target_mode == "delphi2m": + loss, parts = criterion( + logits=logits, + target_events=batch["target_event_seq"], + target_times=batch["target_time_seq"], + current_times=batch["time_seq"], + padding_mask=readout_out.readout_mask, + return_components=True, + ) + else: + loss, parts = criterion( + logits=logits, + target_multi_hot=batch["target_multi_hot"], + target_dt_unique=batch["target_dt_unique"], + readout_mask=readout_out.readout_mask, + return_components=True, + ) + if not torch.isfinite(loss): + raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}") + return loss, parts + + +def run_epoch( + logger: logging.Logger, + args: argparse.Namespace, + model: DeepHealth, + readout, + criterion, + loader: DataLoader, + optimizer: AdamW | None, + device: torch.device, + is_train: bool, +) -> float: + model.train(is_train) + readout.train(is_train) + total = 0.0 + n_batches = 0 + skipped = 0 + parts_sum: Dict[str, float] = {} + desc = "train" if is_train else "val" + + progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True) + for batch_idx, batch in enumerate(progress): + try: + loss, parts = compute_next_step_loss(args, model, readout, criterion, batch, device) + if is_train: + if optimizer is None: + raise ValueError("optimizer is required for training") + optimizer.zero_grad(set_to_none=True) + loss.backward() + if args.grad_clip > 0: + clip_grad_norm_(model.parameters(), args.grad_clip) + optimizer.step() + + total += float(loss.detach().cpu()) + n_batches += 1 + for name, value in parts.items(): + parts_sum[name] = parts_sum.get(name, 0.0) + float(value.detach().cpu()) + avg = total / max(1, n_batches) + postfix = {"loss": f"{float(loss.detach().cpu()):.4f}", "avg": f"{avg:.4f}", "skipped": skipped} + for name, value in parts_sum.items(): + postfix[name] = f"{value / max(1, n_batches):.4f}" + progress.set_postfix(postfix) + except RuntimeError as exc: + if "Loss is not finite" not in str(exc): + raise + skipped += 1 + logger.warning(f"Batch {batch_idx} skipped: {str(exc)[:120]}") + + if skipped: + logger.info(f"Skipped {skipped} batches due to non-finite loss") + return total / max(1, n_batches) if n_batches else float("inf") + + +def build_metadata( + args: argparse.Namespace, + dataset: HealthDataset, + run_name: str, + train_subset, + val_subset, + test_subset, +) -> Dict[str, Any]: + return { + "run_name": run_name, + "dataset_class": "NextStepHealthDataset", + "collate_fn": "next_step_collate_fn", + "model_class": "DeepHealth", + "model_target_mode": "next_token", + "target_mode": args.target_mode, + "dist_mode": "exponential", + "extra_info_types_file": ( + Path(args.extra_info_types_file).name + if args.extra_info_types_file is not None + else None + ), + "extra_info_types": [int(x) for x in dataset.extra_info_types], + "dataset_metadata": { + "vocab_size": int(dataset.vocab_size), + "n_types": int(dataset.n_types), + "n_cont_types": int(dataset.n_cont_types), + "n_categories": int(dataset.n_categories), + "cont_type_ids": [int(x) for x in dataset.cont_type_ids], + "extra_info_types": [int(x) for x in dataset.extra_info_types], + }, + "split_sizes": { + "train": int(len(train_subset)), + "val": int(len(val_subset)), + "test": int(len(test_subset)), + }, + "resolved_readout_name": args.readout_name, + "resolved_loss_name": args.target_mode, + } + + +def main() -> None: + args = parse_args() + set_seed(args.seed) + device = resolve_device(args.device) + configure_torch_for_training(device) + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + run_name = ( + f"{args.time_mode}_exponential_next_token_{args.target_mode}_" + f"gap_{args.no_event_interval_years:g}y_{timestamp}" + ) + run_dir = Path("runs") / run_name + run_dir.mkdir(parents=True, exist_ok=False) + logger = setup_logging(run_dir) + + logger.info(f"Starting next-step training run: {run_name}") + logger.info(f"Device: {device}") + logger.info(f"extra_info_types: {args.extra_info_types or 'all'}") + logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}") + + dataset = HealthDataset( + data_prefix=args.data_prefix, + labels_file=args.labels_file, + no_event_interval_years=args.no_event_interval_years, + include_no_event_in_uts_target=args.include_no_event_in_uts_target, + extra_info_types=args.extra_info_types, + ) + train_subset, val_subset, test_subset = split_dataset( + dataset=dataset, + train_ratio=args.train_ratio, + val_ratio=args.val_ratio, + test_ratio=args.test_ratio, + seed=args.seed, + ) + logger.info( + f"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}" + ) + + train_loader = DataLoader( + train_subset, + batch_size=args.batch_size, + sampler=RandomSampler(train_subset, generator=torch.Generator().manual_seed(args.seed)), + collate_fn=collate_fn, + num_workers=args.num_workers, + pin_memory=device.type == "cuda", + persistent_workers=args.num_workers > 0, + prefetch_factor=2 if args.num_workers > 0 else None, + ) + val_loader = DataLoader( + val_subset, + batch_size=args.batch_size, + shuffle=False, + collate_fn=collate_fn, + num_workers=args.num_workers, + pin_memory=device.type == "cuda", + persistent_workers=args.num_workers > 0, + prefetch_factor=2 if args.num_workers > 0 else None, + ) + test_loader = DataLoader( + test_subset, + batch_size=args.batch_size, + shuffle=False, + collate_fn=collate_fn, + num_workers=args.num_workers, + pin_memory=device.type == "cuda", + persistent_workers=args.num_workers > 0, + prefetch_factor=2 if args.num_workers > 0 else None, + ) + + model = build_model(args, dataset).to(device) + readout = build_next_step_readout(args).to(device) + criterion = build_next_step_loss(args) + optimizer = AdamW( + model.parameters(), + lr=args.base_lr, + betas=tuple(args.betas), + weight_decay=args.weight_decay, + ) + adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128) + + save_config( + args, + run_dir / "train_config.json", + extra=build_metadata(args, dataset, run_name, train_subset, val_subset, test_subset), + ) + + best_val = float("inf") + patience = 0 + history = [] + best_model_path = run_dir / "best_model.pt" + start = time.time() + + for epoch in range(args.max_epochs): + lr = get_lr(epoch, args, adaptive_lr) + set_optimizer_lr(optimizer, lr) + + train_loss = run_epoch(logger, args, model, readout, criterion, train_loader, optimizer, device, True) + with torch.no_grad(): + val_loss = run_epoch(logger, args, model, readout, criterion, val_loader, None, device, False) + + is_best = val_loss < best_val + if is_best: + best_val = val_loss + patience = 0 + save_checkpoint(model, best_model_path) + else: + patience += 1 + + logger.info( + f"Epoch {epoch + 1}/{args.max_epochs} | lr={lr:.6f} | " + f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | " + f"best_val_loss={best_val:.6f} | patience={patience}/{args.patience} | " + f"elapsed={time.time() - start:.1f}s" + ) + history.append({ + "epoch": epoch + 1, + "lr": lr, + "train_loss": train_loss, + "val_loss": val_loss, + "best_val_loss": best_val, + "is_best": int(is_best), + }) + if patience >= args.patience: + logger.info(f"Early stopping triggered at epoch {epoch + 1}") + break + + with (run_dir / "history.json").open("w", encoding="utf-8") as f: + json.dump(history, f, indent=2) + + logger.info("Evaluating best model on next-step test split...") + model.load_state_dict(torch.load(best_model_path, map_location=device)) + with torch.no_grad(): + test_loss = run_epoch(logger, args, model, readout, criterion, test_loader, None, device, False) + logger.info(f"Test loss: {test_loss:.6f}") + logger.info(f"Best checkpoint: {best_model_path}") + + +if __name__ == "__main__": + main() diff --git a/train_util.py b/train_util.py new file mode 100644 index 0000000..b04fd45 --- /dev/null +++ b/train_util.py @@ -0,0 +1,196 @@ +from __future__ import annotations + +import json +import logging +import sys +from pathlib import Path +from typing import Any, Dict, Iterable, Tuple + +import numpy as np +import torch +from torch.optim import AdamW +from torch.utils.data import Subset + +from dataset import AllFutureHealthDataset, HealthDataset +from models import DeepHealth + + +def setup_logging(run_dir: Path) -> logging.Logger: + run_dir.mkdir(parents=True, exist_ok=True) + logger = logging.getLogger("DeepHealth") + logger.setLevel(logging.INFO) + logger.handlers.clear() + + formatter = logging.Formatter( + "%(asctime)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + + console_handler = logging.StreamHandler(sys.stdout) + console_handler.setLevel(logging.INFO) + console_handler.setFormatter(formatter) + logger.addHandler(console_handler) + + file_handler = logging.FileHandler(run_dir / "train.log", mode="w") + file_handler.setLevel(logging.INFO) + file_handler.setFormatter(formatter) + logger.addHandler(file_handler) + + return logger + + +def set_seed(seed: int) -> None: + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + + +def load_extra_info_types_file(path: str) -> list[int]: + file_path = Path(path) + if not file_path.is_file(): + raise FileNotFoundError(f"extra_info_types_file not found: {path}") + + text = file_path.read_text(encoding="utf-8").strip() + if not text: + return [] + + if text.startswith("["): + raw_items = json.loads(text) + if not isinstance(raw_items, list): + raise ValueError("extra_info_types_file JSON must be a list") + else: + raw_items = [] + for line in text.splitlines(): + line = line.split("#", 1)[0].strip() + if line: + raw_items.extend(line.replace(",", " ").replace(";", " ").split()) + + try: + return [int(x) for x in raw_items] + except (TypeError, ValueError) as exc: + raise ValueError(f"Invalid extra info type id in {path}") from exc + + +def configure_torch_for_training(device: torch.device) -> None: + if device.type != "cuda": + return + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + if hasattr(torch, "set_float32_matmul_precision"): + torch.set_float32_matmul_precision("high") + + +def resolve_device(device_arg: str) -> torch.device: + requested = device_arg.strip().lower() + if requested == "cpu": + return torch.device("cpu") + if requested == "cuda": + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + if requested.startswith("cuda:"): + if not torch.cuda.is_available(): + return torch.device("cpu") + index = int(requested.split(":", 1)[1]) + if index < 0 or index >= torch.cuda.device_count(): + raise ValueError(f"Requested CUDA device is out of range: {device_arg}") + return torch.device(f"cuda:{index}") + raise ValueError(f"Unsupported device: {device_arg}") + + +def split_dataset( + dataset: HealthDataset, + train_ratio: float, + val_ratio: float, + test_ratio: float, + seed: int, +) -> Tuple[Subset, Subset, Subset]: + total = train_ratio + val_ratio + test_ratio + if not np.isclose(total, 1.0, atol=1e-6): + raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}") + + indices = np.random.RandomState(seed).permutation(len(dataset)) + n_train = int(len(dataset) * train_ratio) + n_val = int(len(dataset) * val_ratio) + return ( + Subset(dataset, indices[:n_train]), + Subset(dataset, indices[n_train:n_train + n_val]), + Subset(dataset, indices[n_train + n_val:]), + ) + + +def split_all_future_datasets( + train_dataset: AllFutureHealthDataset, + val_dataset: AllFutureHealthDataset, + test_dataset: AllFutureHealthDataset, + train_ratio: float, + val_ratio: float, + test_ratio: float, + seed: int, +) -> Tuple[Subset, Subset, Subset]: + total = train_ratio + val_ratio + test_ratio + if not np.isclose(total, 1.0, atol=1e-6): + raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}") + + patient_indices = np.random.RandomState(seed).permutation(len(train_dataset.patients)) + n_train = int(len(patient_indices) * train_ratio) + n_val = int(len(patient_indices) * val_ratio) + train_patient_idx = patient_indices[:n_train] + val_patient_set = set(int(x) for x in patient_indices[n_train:n_train + n_val]) + test_patient_set = set(int(x) for x in patient_indices[n_train + n_val:]) + + val_query_idx = [ + i for i, (pidx, _t_query) in enumerate(val_dataset.valid_queries) + if int(pidx) in val_patient_set + ] + test_query_idx = [ + i for i, (pidx, _t_query) in enumerate(test_dataset.valid_queries) + if int(pidx) in test_patient_set + ] + if not val_query_idx: + raise ValueError("All-future validation split has no valid query samples.") + if not test_query_idx: + raise ValueError("All-future test split has no valid query samples.") + + return ( + Subset(train_dataset, train_patient_idx), + Subset(val_dataset, np.asarray(val_query_idx, dtype=np.int64)), + Subset(test_dataset, np.asarray(test_query_idx, dtype=np.int64)), + ) + + +def build_optimizer(args: Any, model: DeepHealth) -> AdamW: + return AdamW( + model.parameters(), + lr=args.base_lr, + betas=tuple(args.betas), + weight_decay=args.weight_decay, + ) + + +def set_optimizer_lr(optimizer: AdamW, lr: float) -> None: + for param_group in optimizer.param_groups: + param_group["lr"] = lr + + +def save_checkpoint(model: DeepHealth, checkpoint_path: Path) -> None: + torch.save(model.state_dict(), checkpoint_path) + + +def save_config( + args: Any, + config_path: Path, + extra: Dict[str, Any] | None = None, +) -> None: + config: Dict[str, Any] = {} + for key, value in vars(args).items(): + if isinstance(value, tuple): + config[key] = list(value) + elif isinstance(value, list): + config[key] = value + elif isinstance(value, (int, float, str, bool, type(None))): + config[key] = value + else: + config[key] = str(value) + if extra: + config.update(extra) + config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")