From 517d573d3b4e70d01dd43d0a1a745827a52046dc Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Sat, 27 Jun 2026 15:56:12 +0800 Subject: [PATCH] Reduce attribution GPU synchronization --- ...te_single_disease_mortality_attribution.py | 60 ++++++++++--------- 1 file changed, 32 insertions(+), 28 deletions(-) diff --git a/evaluate_single_disease_mortality_attribution.py b/evaluate_single_disease_mortality_attribution.py index aceb2bc..e254259 100644 --- a/evaluate_single_disease_mortality_attribution.py +++ b/evaluate_single_disease_mortality_attribution.py @@ -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