From 263267f5839b9aea63884908139833251d58820a Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Fri, 26 Jun 2026 11:50:12 +0800 Subject: [PATCH] Use DataLoader for burden readouts --- compute_burden_index_landmarks.py | 498 +++++++++++++++++------------- 1 file changed, 278 insertions(+), 220 deletions(-) diff --git a/compute_burden_index_landmarks.py b/compute_burden_index_landmarks.py index 49055f4..965d156 100644 --- a/compute_burden_index_landmarks.py +++ b/compute_burden_index_landmarks.py @@ -10,6 +10,8 @@ from typing import Any, Iterable import numpy as np import pandas as pd import torch +from torch.nn.utils.rnn import pad_sequence +from torch.utils.data import DataLoader, IterableDataset, get_worker_info from tqdm.auto import tqdm from burden_index import ( @@ -294,117 +296,135 @@ def _config_split_indices( return make_eval_indices(_Sized(), args, cfg) -def _iter_readout_batches( - n: int, - batch_size: int, -) -> Iterable[slice]: - batch_size = max(1, int(batch_size)) - for start in range(0, n, batch_size): - yield slice(start, min(start + batch_size, n)) +class ReadoutJobIterableDataset(IterableDataset): + def __init__( + self, + *, + rows: list[dict[str, Any]], + formed_mode: str, + horizon: float, + ) -> None: + super().__init__() + self.rows = rows + self.formed_mode = str(formed_mode) + self.horizon = float(horizon) + if self.formed_mode not in {"observed", "model_weighted"}: + raise ValueError(f"Unknown formed_mode={self.formed_mode!r}") + + def __iter__(self) -> Iterable[dict[str, torch.Tensor]]: + worker = get_worker_info() + if worker is None: + start, step = 0, 1 + else: + start, step = int(worker.id), int(worker.num_workers) + + for row_idx in range(start, len(self.rows), step): + row = self.rows[row_idx] + if self.formed_mode == "model_weighted": + grid = _build_readout_grid( + event_seq=row["event_seq"], + time_seq=row["time_seq"], + other_type=row["other_type"], + other_time=row["other_time"], + t_query=float(row["t_query"]), + ) + if grid.size > 0: + end_times = np.concatenate( + [grid[1:], np.asarray([row["t_query"]], dtype=np.float32)] + ) + row_deltas = np.maximum(end_times - grid, 0.0).astype(np.float32) + valid = row_deltas > 0 + for query_time, delta in zip(grid[valid].tolist(), row_deltas[valid].tolist()): + yield _make_readout_job(row, row_idx, "formed", query_time, delta) + if self.horizon > 0: + yield _make_readout_job( + row, + row_idx, + "future", + float(row["t_query"]), + self.horizon, + ) + + +def _make_readout_job( + row: dict[str, Any], + row_idx: int, + kind: str, + query_time: float, + delta: float, +) -> dict[str, torch.Tensor]: + return { + "event_seq": torch.from_numpy(np.asarray(row["event_seq"], dtype=np.int64)).long(), + "time_seq": torch.from_numpy(np.asarray(row["time_seq"], dtype=np.float32)).float(), + "sex": torch.tensor(int(row["sex"]), dtype=torch.long), + "other_type": torch.from_numpy(np.asarray(row["other_type"], dtype=np.int64)).long(), + "other_value": torch.from_numpy(np.asarray(row["other_value"], dtype=np.float32)).float(), + "other_value_kind": torch.from_numpy( + np.asarray(row["other_value_kind"], dtype=np.int64) + ).long(), + "other_time": torch.from_numpy(np.asarray(row["other_time"], dtype=np.float32)).float(), + "query_time": torch.tensor(float(query_time), dtype=torch.float32), + "delta": torch.tensor(float(delta), dtype=torch.float32), + "row_idx": torch.tensor(int(row_idx), dtype=torch.long), + "kind": torch.tensor(0 if kind == "formed" else 1, dtype=torch.long), + } + + +def _collate_readout_jobs(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 + ) + return { + "event_seq": event_seq, + "time_seq": time_seq, + "padding_mask": event_seq > PAD_IDX, + "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, + "query_time": torch.stack([x["query_time"] for x in batch]), + "delta": torch.stack([x["delta"] for x in batch]), + "row_idx": torch.stack([x["row_idx"] for x in batch]), + "kind": torch.stack([x["kind"] for x in batch]), + } @torch.inference_mode() -def _query_hidden_jobs( +def _query_hidden_readout_batch( *, ctx: Any, - jobs: list[tuple[dict[str, Any], float]], + batch: dict[str, torch.Tensor], ) -> torch.Tensor: - if not jobs: - return torch.empty(0, ctx.model.n_embd, device=ctx.device) - - batch_size = len(jobs) - max_event_len = max(int(np.asarray(row["event_seq"]).size) for row, _ in jobs) - max_other_len = max(int(np.asarray(row["other_type"]).size) for row, _ in jobs) - - event = np.full((batch_size, max_event_len), PAD_IDX, dtype=np.int64) - time = np.zeros((batch_size, max_event_len), dtype=np.float32) - other_type = np.zeros((batch_size, max_other_len), dtype=np.int64) - other_value = np.zeros((batch_size, max_other_len), dtype=np.float32) - other_value_kind = np.zeros((batch_size, max_other_len), dtype=np.int64) - other_time = np.zeros((batch_size, max_other_len), dtype=np.float32) - sex = np.zeros(batch_size, dtype=np.int64) - query_times = np.zeros(batch_size, dtype=np.float32) - - for i, (row, query_time) in enumerate(jobs): - event_seq = np.asarray(row["event_seq"], dtype=np.int64) - time_seq = np.asarray(row["time_seq"], dtype=np.float32) - other_type_seq = np.asarray(row["other_type"], dtype=np.int64) - other_value_seq = np.asarray(row["other_value"], dtype=np.float32) - other_value_kind_seq = np.asarray(row["other_value_kind"], dtype=np.int64) - other_time_seq = np.asarray(row["other_time"], dtype=np.float32) - - event[i, : event_seq.size] = event_seq - time[i, : time_seq.size] = time_seq - other_type[i, : other_type_seq.size] = other_type_seq - other_value[i, : other_value_seq.size] = other_value_seq - other_value_kind[i, : other_value_kind_seq.size] = other_value_kind_seq - other_time[i, : other_time_seq.size] = other_time_seq - sex[i] = int(row["sex"]) - query_times[i] = np.float32(query_time) - - event_t = torch.from_numpy(event).long().to(ctx.device) + event = batch["event_seq"].long().to(ctx.device, non_blocking=True) return ctx.model( - event_seq=event_t, - time_seq=torch.from_numpy(time).float().to(ctx.device), - sex=torch.from_numpy(sex).long().to(ctx.device), - padding_mask=event_t > PAD_IDX, - t_query=torch.from_numpy(query_times).float().to(ctx.device), - other_type=torch.from_numpy(other_type).long().to(ctx.device), - other_value=torch.from_numpy(other_value).float().to(ctx.device), - other_value_kind=torch.from_numpy(other_value_kind).long().to(ctx.device), - other_time=torch.from_numpy(other_time).float().to(ctx.device), + event_seq=event, + time_seq=batch["time_seq"].float().to(ctx.device, non_blocking=True), + sex=batch["sex"].long().to(ctx.device, non_blocking=True), + padding_mask=event > PAD_IDX, + t_query=batch["query_time"].float().to(ctx.device, non_blocking=True), + other_type=batch["other_type"].long().to(ctx.device, non_blocking=True), + other_value=batch["other_value"].float().to(ctx.device, non_blocking=True), + other_value_kind=batch["other_value_kind"].long().to(ctx.device, non_blocking=True), + other_time=batch["other_time"].float().to(ctx.device, non_blocking=True), target_mode="all_future", ) - -def _build_readout_table( - *, - rows: list[dict[str, Any]], - formed_mode: str, - horizon: float, -) -> dict[str, Any]: - jobs: list[tuple[dict[str, Any], float]] = [] - row_indices: list[int] = [] - kinds: list[str] = [] - deltas: list[float] = [] - - if formed_mode not in {"observed", "model_weighted"}: - raise ValueError(f"Unknown formed_mode={formed_mode!r}") - - for row_idx, row in enumerate(rows): - if formed_mode == "model_weighted": - grid = _build_readout_grid( - event_seq=row["event_seq"], - time_seq=row["time_seq"], - other_type=row["other_type"], - other_time=row["other_time"], - t_query=float(row["t_query"]), - ) - if grid.size > 0: - end_times = np.concatenate( - [grid[1:], np.asarray([row["t_query"]], dtype=np.float32)] - ) - row_deltas = np.maximum(end_times - grid, 0.0).astype(np.float32) - valid = row_deltas > 0 - for query_time, delta in zip(grid[valid].tolist(), row_deltas[valid].tolist()): - jobs.append((row, float(query_time))) - row_indices.append(row_idx) - kinds.append("formed") - deltas.append(float(delta)) - if horizon > 0: - jobs.append((row, float(row["t_query"]))) - row_indices.append(row_idx) - kinds.append("future") - deltas.append(float(horizon)) - - return { - "jobs": jobs, - "row_indices": np.asarray(row_indices, dtype=np.int64), - "kinds": np.asarray(kinds, dtype=object), - "deltas": np.asarray(deltas, dtype=np.float32), - } - - @torch.inference_mode() def _probabilities_from_hidden_torch( *, @@ -442,30 +462,21 @@ def _probabilities_from_hidden_torch( return -torch.expm1(-rate * exposure) - @torch.inference_mode() -def _readout_probabilities( +def _readout_probabilities_from_batch( *, ctx: Any, - readout_table: dict[str, Any], + batch: dict[str, torch.Tensor], union_disease_ids: np.ndarray, - readout_batch_size: int, ) -> torch.Tensor: - jobs = readout_table["jobs"] - if not jobs: - return torch.empty((0, union_disease_ids.size), dtype=torch.float32, device=ctx.device) - - out = torch.empty((len(jobs), union_disease_ids.size), dtype=torch.float32, device=ctx.device) - deltas = np.asarray(readout_table["deltas"], dtype=np.float32) - for slc in _iter_readout_batches(len(jobs), readout_batch_size): - hidden = _query_hidden_jobs(ctx=ctx, jobs=jobs[slc]) - out[slc] = _probabilities_from_hidden_torch( - ctx=ctx, - hidden=hidden, - disease_ids=union_disease_ids, - deltas=deltas[slc], - ).to(dtype=out.dtype) - return out + hidden = _query_hidden_readout_batch(ctx=ctx, batch=batch) + deltas = batch["delta"].detach().cpu().numpy().astype(np.float32, copy=False) + return _probabilities_from_hidden_torch( + ctx=ctx, + hidden=hidden, + disease_ids=union_disease_ids, + deltas=deltas, + ).to(dtype=torch.float32) def _observed_formed_for_rows( @@ -484,94 +495,58 @@ def _observed_formed_for_rows( return formed -def _reduce_readout_table_to_bi_rows( +def _apply_readout_batch_to_accumulators( + *, + batch: dict[str, torch.Tensor], + readout_prob: torch.Tensor, + survival_by_row: torch.Tensor | None, + future_prob_by_row: torch.Tensor, + ctx: Any, +) -> None: + if readout_prob.numel() == 0: + return + + row_indices = batch["row_idx"].long().to(ctx.device, non_blocking=True) + kinds = batch["kind"].long().to(ctx.device, non_blocking=True) + kind_is_formed = kinds == 0 + kind_is_future = kinds == 1 + + if survival_by_row is not None and bool(kind_is_formed.any().item()): + formed_rows = row_indices[kind_is_formed] + formed_survival = 1.0 - readout_prob[kind_is_formed].clamp(0.0, 1.0) + if hasattr(survival_by_row, "scatter_reduce_"): + survival_by_row.scatter_reduce_( + dim=0, + index=formed_rows[:, None].expand_as(formed_survival), + src=formed_survival, + reduce="prod", + include_self=True, + ) + else: + for job_idx in torch.nonzero(kind_is_formed, as_tuple=False).flatten().tolist(): + survival_by_row[int(row_indices[job_idx].item())] *= ( + 1.0 - readout_prob[job_idx].clamp(0.0, 1.0) + ) + if bool(kind_is_future.any().item()): + future_rows = row_indices[kind_is_future] + future_prob_by_row[future_rows] = readout_prob[kind_is_future] + + +def _project_bi_rows( *, rows: list[dict[str, Any]], horizon: float, matrices: list[dict[str, Any]], - union_disease_ids: np.ndarray, formed_mode: str, - readout_table: dict[str, Any], - readout_prob: torch.Tensor, + formed_by_row: torch.Tensor, + future_prob_by_row: torch.Tensor, ctx: Any, ) -> list[dict[str, Any]]: - if formed_mode == "observed": - formed_by_row = torch.as_tensor( - _observed_formed_for_rows( - rows=rows, - union_disease_ids=union_disease_ids, - ), - dtype=readout_prob.dtype, - device=ctx.device, - ) - elif formed_mode == "model_weighted": - formed_by_row = torch.zeros( - (len(rows), union_disease_ids.size), - dtype=readout_prob.dtype, - device=ctx.device, - ) - else: - raise ValueError(f"Unknown formed_mode={formed_mode!r}") - - future_prob_by_row = torch.zeros( - (len(rows), union_disease_ids.size), - dtype=readout_prob.dtype, - device=ctx.device, - ) - row_indices = torch.as_tensor( - np.asarray(readout_table["row_indices"], dtype=np.int64), - dtype=torch.long, - device=ctx.device, - ) - kinds = np.asarray(readout_table["kinds"], dtype=object) - if formed_mode == "model_weighted": - survival_by_row = torch.ones( - (len(rows), union_disease_ids.size), - dtype=readout_prob.dtype, - device=ctx.device, - ) - else: - survival_by_row = None - - if readout_prob.numel() > 0: - kind_is_formed = torch.as_tensor( - np.asarray(kinds == "formed", dtype=np.bool_), - dtype=torch.bool, - device=ctx.device, - ) - kind_is_future = torch.as_tensor( - np.asarray(kinds == "future", dtype=np.bool_), - dtype=torch.bool, - device=ctx.device, - ) - if survival_by_row is not None and bool(kind_is_formed.any().item()): - formed_rows = row_indices[kind_is_formed] - formed_survival = 1.0 - readout_prob[kind_is_formed].clamp(0.0, 1.0) - if hasattr(survival_by_row, "scatter_reduce_"): - survival_by_row.scatter_reduce_( - dim=0, - index=formed_rows[:, None].expand_as(formed_survival), - src=formed_survival, - reduce="prod", - include_self=True, - ) - else: - for job_idx in torch.nonzero(kind_is_formed, as_tuple=False).flatten().tolist(): - survival_by_row[int(row_indices[job_idx].item())] *= ( - 1.0 - readout_prob[job_idx].clamp(0.0, 1.0) - ) - if bool(kind_is_future.any().item()): - future_rows = row_indices[kind_is_future] - future_prob_by_row[future_rows] = readout_prob[kind_is_future] - - if survival_by_row is not None: - formed_by_row = 1.0 - survival_by_row - disease_future_by_row = (1.0 - formed_by_row) * future_prob_by_row disease_total_by_row = formed_by_row + disease_future_by_row projected: list[dict[str, Any]] = [] for matrix in matrices: - A = torch.as_tensor(matrix["A_union"], dtype=readout_prob.dtype, device=ctx.device) + A = torch.as_tensor(matrix["A_union"], dtype=formed_by_row.dtype, device=ctx.device) projected.append( { "matrix": matrix, @@ -611,7 +586,7 @@ def _reduce_readout_table_to_bi_rows( return out -def _compute_bi_from_readout_table( +def _compute_bi_from_streamed_readouts( *, rows: list[dict[str, Any]], horizon: float, @@ -619,46 +594,117 @@ def _compute_bi_from_readout_table( union_disease_ids: np.ndarray, formed_mode: str, readout_batch_size: int, + num_workers: int, ctx: Any, + log_prefix: str | None = None, ) -> tuple[list[dict[str, Any]], int, dict[str, float]]: horizon = float(horizon) if horizon < 0: raise ValueError(f"horizon must be non-negative, got {horizon}") - t0 = time.perf_counter() - readout_table = _build_readout_table( + + dtype = torch.float32 + if formed_mode == "observed": + formed_by_row = torch.as_tensor( + _observed_formed_for_rows( + rows=rows, + union_disease_ids=union_disease_ids, + ), + dtype=dtype, + device=ctx.device, + ) + survival_by_row = None + elif formed_mode == "model_weighted": + survival_by_row = torch.ones( + (len(rows), union_disease_ids.size), + dtype=dtype, + device=ctx.device, + ) + formed_by_row = None + else: + raise ValueError(f"Unknown formed_mode={formed_mode!r}") + future_prob_by_row = torch.zeros( + (len(rows), union_disease_ids.size), + dtype=dtype, + device=ctx.device, + ) + + readout_jobs = 0 + n_batches = 0 + build_readout_sec = 0.0 + forward_sec = 0.0 + reduce_sec = 0.0 + t_loader0 = time.perf_counter() + readout_dataset = ReadoutJobIterableDataset( rows=rows, formed_mode=formed_mode, horizon=horizon, ) - t1 = time.perf_counter() - readout_prob = _readout_probabilities( - ctx=ctx, - readout_table=readout_table, - union_disease_ids=union_disease_ids, - readout_batch_size=readout_batch_size, + readout_loader = DataLoader( + readout_dataset, + batch_size=max(1, int(readout_batch_size)), + collate_fn=_collate_readout_jobs, + num_workers=max(0, int(num_workers)), + pin_memory=ctx.device.type == "cuda", + persistent_workers=int(num_workers) > 0, + prefetch_factor=2 if int(num_workers) > 0 else None, ) - if ctx.device.type == "cuda": - torch.cuda.synchronize(ctx.device) - t2 = time.perf_counter() - rows_out = _reduce_readout_table_to_bi_rows( + build_readout_sec += time.perf_counter() - t_loader0 + + for batch in readout_loader: + t_forward0 = time.perf_counter() + readout_prob = _readout_probabilities_from_batch( + ctx=ctx, + batch=batch, + union_disease_ids=union_disease_ids, + ) + if ctx.device.type == "cuda": + torch.cuda.synchronize(ctx.device) + forward_sec += time.perf_counter() - t_forward0 + + t_reduce0 = time.perf_counter() + _apply_readout_batch_to_accumulators( + batch=batch, + readout_prob=readout_prob, + survival_by_row=survival_by_row, + future_prob_by_row=future_prob_by_row, + ctx=ctx, + ) + if ctx.device.type == "cuda": + torch.cuda.synchronize(ctx.device) + reduce_sec += time.perf_counter() - t_reduce0 + + n_batches += 1 + readout_jobs += int(batch["row_idx"].numel()) + if log_prefix and (n_batches == 1 or n_batches % 50 == 0): + print( + f"{log_prefix} processed {readout_jobs} readout jobs " + f"in {n_batches} batches", + flush=True, + ) + + if survival_by_row is not None: + formed_by_row = 1.0 - survival_by_row + assert formed_by_row is not None + + t_project0 = time.perf_counter() + rows_out = _project_bi_rows( rows=rows, horizon=horizon, matrices=matrices, - union_disease_ids=union_disease_ids, formed_mode=formed_mode, - readout_table=readout_table, - readout_prob=readout_prob, + formed_by_row=formed_by_row, + future_prob_by_row=future_prob_by_row, ctx=ctx, ) if ctx.device.type == "cuda": torch.cuda.synchronize(ctx.device) - t3 = time.perf_counter() + reduce_sec += time.perf_counter() - t_project0 timings = { - "build_readout_sec": t1 - t0, - "forward_sec": t2 - t1, - "reduce_sec": t3 - t2, + "build_readout_sec": build_readout_sec, + "forward_sec": forward_sec, + "reduce_sec": reduce_sec, } - return rows_out, len(readout_table["jobs"]), timings + return rows_out, readout_jobs, timings def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]: @@ -681,14 +727,16 @@ def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]: f"{time.perf_counter() - materialize_start:.2f}s", flush=True, ) - out, readout_jobs, timings = _compute_bi_from_readout_table( + out, readout_jobs, timings = _compute_bi_from_streamed_readouts( rows=rows, horizon=payload["horizon"], matrices=payload["matrices"], union_disease_ids=payload["union_disease_ids"], formed_mode=payload["formed_mode"], readout_batch_size=int(payload["readout_batch_size"]), + num_workers=int(payload["num_workers"]), ctx=ctx, + log_prefix=f"[BI worker {device}]", ) print( f"[BI worker {device}] done: readout_jobs={readout_jobs}, " @@ -831,6 +879,12 @@ def main() -> None: "Increase this to improve GPU utilization if memory allows." ), ) + parser.add_argument( + "--num_workers", + type=int, + default=4, + help="DataLoader workers per GPU process for readout job generation.", + ) parser.add_argument("--device", type=str, default=None) parser.add_argument( "--devices", @@ -933,14 +987,16 @@ def main() -> None: flush=True, ) if len(row_chunks) == 1: - all_rows, total_readout_jobs, timings = _compute_bi_from_readout_table( + all_rows, total_readout_jobs, timings = _compute_bi_from_streamed_readouts( rows=rows, horizon=horizon, matrices=matrices, union_disease_ids=union_disease_ids, formed_mode=args.formed_mode, readout_batch_size=int(args.readout_batch_size), + num_workers=int(args.num_workers), ctx=ctx, + log_prefix="[BI main]", ) for key, value in timings.items(): total_timings[key] += float(value) @@ -957,6 +1013,7 @@ def main() -> None: "matrices": matrices, "union_disease_ids": union_disease_ids, "readout_batch_size": int(args.readout_batch_size), + "num_workers": int(args.num_workers), "formed_mode": args.formed_mode, } for device, chunk_rows in row_chunks @@ -988,6 +1045,7 @@ def main() -> None: print(f"Landmark rows: {len(rows)}") print(f"Readout jobs: {total_readout_jobs}") print(f"Readout batch size per worker: {int(args.readout_batch_size)}") + print(f"DataLoader workers per GPU process: {int(args.num_workers)}") print(f"Multiprocessing start method: {args.mp_start_method}") print( "Timing seconds: "