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

Keep burden reduction on GPU

Browse Source
This commit is contained in:
Jiarui Li
2026-06-26 11:27:44 +08:00
parent 63df61e1dd
commit c5ecbb79f2
1 changed files with 216 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

257
compute_burden_index_landmarks.py
View File

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