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

Reduce attribution GPU synchronization

Browse Source
This commit is contained in:
Jiarui Li
2026-06-27 15:56:12 +08:00
parent 20686c128f
commit 517d573d3b
1 changed files with 32 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
evaluate_single_disease_mortality_attribution.py
View File

@@ -555,7 +555,7 @@ def main() -> None:
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
attribution_batch_size = int(
cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 8, batch_size))
cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 32, 4096))
)
if attribution_batch_size <= 0:
raise ValueError("attribution_batch_size must be positive")
@@ -613,15 +613,20 @@ def main() -> None:
row_base_cache: dict[int, dict[str, Any]] = {}
pending_batch_chunks: list[Dict[str, torch.Tensor]] = []
pending_meta_chunks: list[list[dict[str, Any]]] = []
pending_orig_risk_chunks: list[torch.Tensor] = []
pending_orig_hazard_chunks: list[torch.Tensor] = []
pending_n = 0
def flush_pending() -> None:
nonlocal written_rows, shard_index, pending_batch_chunks, pending_meta_chunks, pending_n
nonlocal written_rows, shard_index, pending_batch_chunks, pending_meta_chunks
nonlocal pending_orig_risk_chunks, pending_orig_hazard_chunks, pending_n
if pending_n == 0:
return
ablated_batch = concat_padded_tensor_batches(pending_batch_chunks)
meta_rows = [row for chunk in pending_meta_chunks for row in chunk]
orig_risk = torch.cat(pending_orig_risk_chunks, dim=0).to(device)
orig_hazard = torch.cat(pending_orig_hazard_chunks, dim=0).to(device)
with torch.no_grad():
ablated_risk = death_risk_for_batch(
model=model,
@@ -634,35 +639,33 @@ def main() -> None:
tau=tau,
)
ablated_hazard = mortality_hazard_from_risk(ablated_risk)
orig_risk = torch.as_tensor(
[row.pop("_death_risk") for row in meta_rows],
dtype=ablated_risk.dtype,
device=ablated_risk.device,
)
orig_hazard = torch.as_tensor(
[row.pop("_death_hazard") for row in meta_rows],
dtype=ablated_hazard.dtype,
device=ablated_hazard.device,
)
attr_prob = orig_risk - ablated_risk
attr_hazard = orig_hazard - ablated_hazard
ratio_prob = safe_ratio(orig_risk, ablated_risk, eps=float(args.ratio_eps))
ratio_hazard = safe_ratio(orig_hazard, ablated_hazard, eps=float(args.ratio_eps))
value_block = torch.stack(
[
orig_risk,
orig_hazard,
ablated_risk,
ablated_hazard,
attr_prob,
attr_hazard,
ratio_prob,
ratio_hazard,
],
dim=1,
).detach().cpu().numpy()
for i, row in enumerate(meta_rows):
row["death_risk"] = float(orig_risk[i].detach().cpu())
row["death_hazard"] = float(orig_hazard[i].detach().cpu())
row["ablated_death_risk"] = float(ablated_risk[i].detach().cpu())
row["ablated_death_hazard"] = float(ablated_hazard[i].detach().cpu())
row["mortality_attribution_probability"] = float(attr_prob[i].detach().cpu())
row["mortality_attribution_hazard"] = float(attr_hazard[i].detach().cpu())
row["mortality_attribution_probability_ratio"] = float(
ratio_prob[i].detach().cpu()
)
row["mortality_attribution_hazard_ratio"] = float(
ratio_hazard[i].detach().cpu()
)
row["death_risk"] = float(value_block[i, 0])
row["death_hazard"] = float(value_block[i, 1])
row["ablated_death_risk"] = float(value_block[i, 2])
row["ablated_death_hazard"] = float(value_block[i, 3])
row["mortality_attribution_probability"] = float(value_block[i, 4])
row["mortality_attribution_hazard"] = float(value_block[i, 5])
row["mortality_attribution_probability_ratio"] = float(value_block[i, 6])
row["mortality_attribution_hazard_ratio"] = float(value_block[i, 7])
table = pd.DataFrame(meta_rows).reindex(columns=OUTPUT_COLUMNS)
update_summary_accumulator(summary_accumulator, table)
@@ -674,6 +677,8 @@ def main() -> None:
written_rows += len(meta_rows)
pending_batch_chunks = []
pending_meta_chunks = []
pending_orig_risk_chunks = []
pending_orig_hazard_chunks = []
pending_n = 0
def get_row_base(row_idx: int) -> dict[str, Any]:
@@ -779,7 +784,6 @@ def main() -> None:
f"{disease_token} for row {row_idx}, but cached history has count 0"
)
orig_row = int(local_rows[local_pos].detach().cpu().item())
meta_chunk.append(
{
"patient_id": row_base["patient_id"],
@@ -803,14 +807,14 @@ def main() -> None:
"history_count__selected_organ_system": int(
group_counts.get(str(disease_meta.get("organ_system", "")), 0)
),
"_death_risk": float(death_risk_tensor[orig_row].detach().cpu()),
"_death_hazard": float(death_hazard_tensor[orig_row].detach().cpu()),
}
)
if meta_chunk:
pending_batch_chunks.append(ablated_chunk)
pending_meta_chunks.append(meta_chunk)
pending_orig_risk_chunks.append(death_risk_tensor[local_rows].detach())
pending_orig_hazard_chunks.append(death_hazard_tensor[local_rows].detach())
pending_n += len(meta_chunk)
pair_offset = pair_stop