From c5ecbb79f213143696b99b35449862fc6c36bbbf Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Fri, 26 Jun 2026 11:27:44 +0800 Subject: [PATCH] Keep burden reduction on GPU --- compute_burden_index_landmarks.py | 257 +++++++++++++++++++++++++----- 1 file changed, 216 insertions(+), 41 deletions(-) diff --git a/compute_burden_index_landmarks.py b/compute_burden_index_landmarks.py index 4599e77..0ea1a31 100644 --- a/compute_burden_index_landmarks.py +++ b/compute_burden_index_landmarks.py @@ -2,6 +2,7 @@ from __future__ import annotations import argparse import multiprocessing as mp +import time from concurrent.futures import ProcessPoolExecutor, as_completed from pathlib import Path from typing import Any, Iterable @@ -14,7 +15,6 @@ from tqdm.auto import tqdm from burden_index import ( _build_readout_grid, _observed_formed_burden, - _probabilities_from_hidden, load_burden_context, ) from evaluate_auc_v2 import ( @@ -361,6 +361,44 @@ def _build_readout_table( } +@torch.inference_mode() +def _probabilities_from_hidden_torch( + *, + ctx: Any, + hidden: torch.Tensor, + disease_ids: np.ndarray, + deltas: np.ndarray, +) -> torch.Tensor: + if hidden.ndim != 2: + raise ValueError(f"hidden must have shape (N, H), got {tuple(hidden.shape)}") + if deltas.ndim != 1 or deltas.size != hidden.shape[0]: + raise ValueError( + "deltas must be 1D with the same length as hidden rows, got " + f"{deltas.shape} vs {tuple(hidden.shape)}" + ) + + ids = torch.as_tensor(disease_ids, dtype=torch.long, device=ctx.device) + logits = ctx.model.calc_risk(hidden)[:, ids] + rate = torch.nn.functional.softplus(logits).clamp_min(1e-8) + delta_t = torch.as_tensor(deltas, dtype=rate.dtype, device=ctx.device).clamp_min(0) + + if ctx.dist_mode == "weibull": + rho = ctx.model.calc_weibull_rho(hidden)[:, ids] + exposure = torch.pow(delta_t[:, None], rho) + elif ctx.dist_mode == "mixed": + exposure = delta_t[:, None].expand_as(rate) + death_idx = int(getattr(ctx.model, "death_idx", getattr(ctx.model, "vocab_size", 0) - 1)) + death_cols = [j for j, token in enumerate(disease_ids.tolist()) if int(token) == death_idx] + if death_cols: + death_rho = ctx.model.calc_death_rho(hidden) + for col in death_cols: + exposure[:, int(col)] = torch.pow(delta_t, death_rho) + else: + exposure = delta_t[:, None].expand_as(rate) + + return -torch.expm1(-rate * exposure) + + @torch.inference_mode() def _readout_probabilities( *, @@ -368,21 +406,21 @@ def _readout_probabilities( readout_table: dict[str, Any], union_disease_ids: np.ndarray, readout_batch_size: int, -) -> np.ndarray: +) -> torch.Tensor: jobs = readout_table["jobs"] if not jobs: - return np.zeros((0, union_disease_ids.size), dtype=np.float64) + return torch.empty((0, union_disease_ids.size), dtype=torch.float32, device=ctx.device) - out = np.empty((len(jobs), union_disease_ids.size), dtype=np.float64) + 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( + out[slc] = _probabilities_from_hidden_torch( ctx=ctx, hidden=hidden, disease_ids=union_disease_ids, deltas=deltas[slc], - ) + ).to(dtype=out.dtype) return out @@ -410,49 +448,102 @@ def _reduce_readout_table_to_bi_rows( union_disease_ids: np.ndarray, formed_mode: str, readout_table: dict[str, Any], - readout_prob: np.ndarray, + readout_prob: torch.Tensor, + ctx: Any, ) -> list[dict[str, Any]]: if formed_mode == "observed": - formed_by_row = _observed_formed_for_rows( - rows=rows, - union_disease_ids=union_disease_ids, + 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 = np.zeros((len(rows), union_disease_ids.size), dtype=np.float64) + 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 = np.zeros((len(rows), union_disease_ids.size), dtype=np.float64) - row_indices = np.asarray(readout_table["row_indices"], dtype=np.int64) + 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 = np.ones((len(rows), union_disease_ids.size), dtype=np.float64) + survival_by_row = torch.ones( + (len(rows), union_disease_ids.size), + dtype=readout_prob.dtype, + device=ctx.device, + ) else: survival_by_row = None - for job_idx, row_idx in enumerate(row_indices.tolist()): - kind = str(kinds[job_idx]) - if kind == "formed" and survival_by_row is not None: - survival_by_row[int(row_idx)] *= 1.0 - np.clip(readout_prob[job_idx], 0.0, 1.0) - elif kind == "future": - future_prob_by_row[int(row_idx)] = readout_prob[job_idx] + 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) + projected.append( + { + "matrix": matrix, + "historical": torch.matmul(formed_by_row, A.T).detach().cpu().numpy(), + "future": torch.matmul(disease_future_by_row, A.T).detach().cpu().numpy(), + "total": torch.matmul(disease_total_by_row, A.T).detach().cpu().numpy(), + } + ) + out: list[dict[str, Any]] = [] for row_idx, row in enumerate(rows): - formed = formed_by_row[row_idx] - historical_by_matrix = { - matrix["burden_type"]: matrix["A_union"] @ formed - for matrix in matrices - } - disease_future = (1.0 - formed) * future_prob_by_row[row_idx] - disease_total = formed + disease_future - for matrix in matrices: - historical = historical_by_matrix[matrix["burden_type"]] - future = matrix["A_union"] @ disease_future - total = matrix["A_union"] @ disease_total + for item in projected: + matrix = item["matrix"] + historical = item["historical"][row_idx] + future = item["future"][row_idx] + total = item["total"][row_idx] for dim_idx, meta in matrix["category_meta"].iterrows(): out.append( { @@ -485,21 +576,26 @@ def _compute_bi_from_readout_table( formed_mode: str, readout_batch_size: int, ctx: Any, -) -> tuple[list[dict[str, Any]], int]: +) -> 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( 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, ) + if ctx.device.type == "cuda": + torch.cuda.synchronize(ctx.device) + t2 = time.perf_counter() rows_out = _reduce_readout_table_to_bi_rows( rows=rows, horizon=horizon, @@ -508,15 +604,24 @@ def _compute_bi_from_readout_table( formed_mode=formed_mode, readout_table=readout_table, readout_prob=readout_prob, + ctx=ctx, ) - return rows_out, len(readout_table["jobs"]) + if ctx.device.type == "cuda": + torch.cuda.synchronize(ctx.device) + t3 = time.perf_counter() + timings = { + "build_readout_sec": t1 - t0, + "forward_sec": t2 - t1, + "reduce_sec": t3 - t2, + } + return rows_out, len(readout_table["jobs"]), timings def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]: device = payload["device"] run_path = Path(payload["run_path"]) ctx = load_burden_context(run_path, device=device) - out, readout_jobs = _compute_bi_from_readout_table( + out, readout_jobs, timings = _compute_bi_from_readout_table( rows=payload["rows"], horizon=payload["horizon"], matrices=payload["matrices"], @@ -525,7 +630,7 @@ def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]: readout_batch_size=int(payload["readout_batch_size"]), ctx=ctx, ) - return {"rows": out, "readout_jobs": readout_jobs} + return {"rows": out, "readout_jobs": readout_jobs, "timings": timings} def _attach_union_projection( @@ -561,18 +666,64 @@ def _attach_union_projection( def _split_rows_for_devices( rows: list[dict[str, Any]], devices: list[str | None], + *, + formed_mode: str, + horizon: float, ) -> list[tuple[str | None, list[dict[str, Any]]]]: if len(devices) <= 1: return [(devices[0], rows)] - index_chunks = np.array_split(np.arange(len(rows)), len(devices)) + + buckets: list[list[dict[str, Any]]] = [[] for _ in devices] + loads = np.zeros(len(devices), dtype=np.int64) + weighted_rows = sorted( + rows, + key=lambda row: _estimate_readout_jobs_for_row( + row, + formed_mode=formed_mode, + horizon=horizon, + ), + reverse=True, + ) + for row in weighted_rows: + bucket_idx = int(np.argmin(loads)) + buckets[bucket_idx].append(row) + loads[bucket_idx] += _estimate_readout_jobs_for_row( + row, + formed_mode=formed_mode, + horizon=horizon, + ) + chunks: list[tuple[str | None, list[dict[str, Any]]]] = [] - for device, idx in zip(devices, index_chunks): - if idx.size == 0: + for device, bucket in zip(devices, buckets): + if not bucket: continue - chunks.append((device, [rows[int(i)] for i in idx.tolist()])) + chunks.append((device, bucket)) return chunks +def _estimate_readout_jobs_for_row( + row: dict[str, Any], + *, + formed_mode: str, + horizon: float, +) -> int: + n_jobs = 0 + 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)]) + n_jobs += int(np.sum(np.maximum(end_times - grid, 0.0) > 0)) + if horizon > 0: + n_jobs += 1 + return max(n_jobs, 1) + + def main() -> None: parser = argparse.ArgumentParser( description="Compute DeepHealth Burden Indices at landmark ages." @@ -682,9 +833,20 @@ def main() -> None: all_rows: list[dict[str, Any]] = [] total_readout_jobs = 0 - row_chunks = _split_rows_for_devices(rows, devices) + total_timings = { + "build_readout_sec": 0.0, + "forward_sec": 0.0, + "reduce_sec": 0.0, + "write_csv_sec": 0.0, + } + row_chunks = _split_rows_for_devices( + rows, + devices, + formed_mode=args.formed_mode, + horizon=horizon, + ) if len(row_chunks) == 1: - all_rows, total_readout_jobs = _compute_bi_from_readout_table( + all_rows, total_readout_jobs, timings = _compute_bi_from_readout_table( rows=rows, horizon=horizon, matrices=matrices, @@ -693,6 +855,8 @@ def main() -> None: readout_batch_size=int(args.readout_batch_size), ctx=ctx, ) + for key, value in timings.items(): + total_timings[key] += float(value) else: # The main-process context is only needed to build the dataset and rows. # Workers load their own model copy on the assigned device. @@ -724,15 +888,26 @@ def main() -> None: result = future.result() all_rows.extend(result["rows"]) total_readout_jobs += int(result["readout_jobs"]) + for key, value in result["timings"].items(): + total_timings[key] += float(value) + write_start = time.perf_counter() out_df = pd.DataFrame(all_rows) out_df.to_csv(output_path, index=False) + total_timings["write_csv_sec"] = time.perf_counter() - write_start print(f"Run path: {run_path}") print(f"Eval split: {eval_split}") print(f"Horizon: {horizon:g}") 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( + "Timing seconds: " + f"build_readout={total_timings['build_readout_sec']:.2f}, " + f"forward={total_timings['forward_sec']:.2f}, " + f"reduce={total_timings['reduce_sec']:.2f}, " + f"write_csv={total_timings['write_csv_sec']:.2f}" + ) print(f"Devices: {', '.join(str(d) for d, _ in row_chunks)}") for matrix in matrices: print(