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

Add training scripts for all-future and next-step supervision with DeepHealth

Browse Source 
- 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.
This commit is contained in:
Jiarui Li
2026-06-13 11:42:04 +08:00
parent 034d8065a7
commit 46a3dfe628
12 changed files with 1927 additions and 1273 deletions
Download Patch File Download Diff File Expand all files Collapse all files

97
dataset.py
View File

@@ -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: