From 20686c128fca3b29a591843c84bd964469b8bf05 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Sat, 27 Jun 2026 15:51:12 +0800 Subject: [PATCH] Vectorize per-disease attribution batches --- ...te_single_disease_mortality_attribution.py | 224 ++++++++++++------ 1 file changed, 151 insertions(+), 73 deletions(-) diff --git a/evaluate_single_disease_mortality_attribution.py b/evaluate_single_disease_mortality_attribution.py index 0104198..aceb2bc 100644 --- a/evaluate_single_disease_mortality_attribution.py +++ b/evaluate_single_disease_mortality_attribution.py @@ -35,7 +35,6 @@ from evaluate_event_free_survival import ( IndexedLandmarkDataset, LandmarkDataset, build_first_occurrence_maps_for_landmarks, - build_group_ablated_slice, collate_indexed_landmark_fn, death_risk_for_batch, historical_counts_by_group, @@ -47,7 +46,7 @@ from evaluate_event_free_survival import ( mortality_hazard_from_risk, ) from future_risk import death_risk_from_probabilities, probabilities_from_logits -from targets import PAD_IDX +from targets import CHECKUP_IDX, PAD_IDX OUTPUT_COLUMNS = [ @@ -239,6 +238,76 @@ def concat_padded_tensor_batches(chunks: list[Dict[str, torch.Tensor]]) -> Dict[ return out +def build_disease_ablated_slice( + batch: Dict[str, torch.Tensor], + row_indices: torch.Tensor, + token_ids: torch.Tensor, +) -> Dict[str, torch.Tensor]: + """Build an ablated slice for aligned (row, disease_token) pairs.""" + event_seq = batch["event_seq"] + row_indices = row_indices.to(device=event_seq.device, dtype=torch.long) + token_ids = token_ids.to(device=event_seq.device, dtype=event_seq.dtype) + + out: Dict[str, torch.Tensor] = {} + out["event_seq"] = event_seq[row_indices].clone() + out["time_seq"] = batch["time_seq"][row_indices] + out["readout_mask"] = batch["readout_mask"][row_indices].clone() + out["padding_mask"] = batch["padding_mask"][row_indices].bool().clone() + out["landmark_pos"] = batch["landmark_pos"][row_indices].clone() + + seq_len = int(event_seq.shape[1]) + positions = torch.arange(seq_len, device=event_seq.device)[None, :] + remove = (out["event_seq"] == token_ids[:, None]) & out["padding_mask"] + out["event_seq"] = torch.where( + remove, + torch.full_like(out["event_seq"], PAD_IDX), + out["event_seq"], + ) + out["padding_mask"] &= ~remove + out["readout_mask"] &= ~remove + + has_valid = out["padding_mask"].any(dim=1) + if not bool(has_valid.all().item()): + empty_rows = torch.nonzero(~has_valid, as_tuple=False).flatten() + out["event_seq"][empty_rows, 0] = CHECKUP_IDX + out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to( + dtype=out["time_seq"].dtype + ) + out["padding_mask"][empty_rows, 0] = True + out["readout_mask"][empty_rows, 0] = True + out["landmark_pos"][empty_rows] = 0 + + has_readout = out["readout_mask"].any(dim=1) + if not bool(has_readout.all().item()): + rows = torch.nonzero(~has_readout, as_tuple=False).flatten() + local_valid = out["padding_mask"][rows] + last_pos = torch.where( + local_valid, + positions.expand(local_valid.shape[0], -1), + torch.zeros_like(positions.expand(local_valid.shape[0], -1)), + ).amax(dim=1) + out["readout_mask"][rows] = False + out["readout_mask"][rows, last_pos] = True + out["landmark_pos"][rows] = last_pos.to(dtype=out["landmark_pos"].dtype) + + repeated_keys = ( + "sex", + "other_type", + "other_value", + "other_value_kind", + "other_time", + "t_query", + "patient_id", + "landmark_age", + "followup_end_time", + "death_time", + "row_idx", + ) + for key in repeated_keys: + out[key] = batch[key][row_indices] + return out + + def load_disease_metadata( mapping_path: Path, *, @@ -486,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 * 4, batch_size)) + cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 8, batch_size)) ) if attribution_batch_size <= 0: raise ValueError("attribution_batch_size must be positive") @@ -616,6 +685,7 @@ def main() -> None: dataset_index = int(meta["dataset_index"]) sample = dataset.samples[dataset_index] hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64) + unique_tokens, token_counts = np.unique(hist_tokens, return_counts=True) total_count, group_counts = historical_counts_by_group( hist_tokens, death_idx=death_idx, @@ -632,15 +702,23 @@ def main() -> None: "followup_end_time": float(meta["followup_end_time"]), "history_disease_count": int(total_count), "_hist_tokens": hist_tokens, + "_token_counts": { + int(token): int(count) + for token, count in zip(unique_tokens.tolist(), token_counts.tolist()) + }, "_group_counts": group_counts, } row_base_cache[row_idx] = cached return cached for batch in tqdm(loader, desc="Per-disease mortality attribution", dynamic_ncols=True): + batch_dev = { + k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v) + for k, v in batch.items() + } hidden = infer_landmark_hidden( model=model, - batch=batch, + batch=batch_dev, device=device, model_target_mode=model_target_mode, readout_name=readout_name, @@ -660,84 +738,84 @@ def main() -> None: death_risk_tensor = death_risk_from_probabilities(probabilities) death_hazard_tensor = mortality_hazard_from_risk(death_risk_tensor) occurred = make_occurred_mask( - batch["event_seq"].to(device), + batch_dev["event_seq"], vocab_size=int(dataset.vocab_size), device=device, ) - active_token_mask = occurred[:, scanned_disease_tensor].any(dim=0) - batch_disease_tokens = scanned_disease_tensor[ - active_token_mask - ].detach().cpu().tolist() - if not batch_disease_tokens: + pair_indices = torch.nonzero( + occurred[:, scanned_disease_tensor], + as_tuple=False, + ) + if pair_indices.numel() == 0: continue - for disease_token in batch_disease_tokens: - disease_meta = metadata[disease_token] - active_rows = torch.nonzero( - occurred[:, disease_token].to(dtype=torch.bool), - as_tuple=False, - ).flatten() - if active_rows.numel() == 0: - continue + pair_offset = 0 + while pair_offset < int(pair_indices.shape[0]): + capacity = int(attribution_batch_size) - pending_n + pair_stop = min(int(pair_indices.shape[0]), pair_offset + capacity) + pair_chunk = pair_indices[pair_offset:pair_stop] + local_rows = pair_chunk[:, 0].long() + disease_token_ids = scanned_disease_tensor[pair_chunk[:, 1]].long() + ablated_chunk = build_disease_ablated_slice( + batch=batch_dev, + row_indices=local_rows, + token_ids=disease_token_ids, + ) - row_offset = 0 - while row_offset < int(active_rows.numel()): - capacity = int(attribution_batch_size) - pending_n - row_stop = min(int(active_rows.numel()), row_offset + capacity) - row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device) - ablated_chunk = build_group_ablated_slice( - batch=batch, - token_ids=[disease_token], - row_indices=row_indices, - ) - - meta_chunk: list[dict[str, Any]] = [] - row_ids = batch["row_idx"][row_indices.cpu()].cpu().numpy().astype(np.int64) - for local_pos, row_idx in enumerate(row_ids.tolist()): - row_base = get_row_base(int(row_idx)) - hist_tokens = row_base["_hist_tokens"] - group_counts = row_base["_group_counts"] - disease_history_count = int((hist_tokens == int(disease_token)).sum()) - if disease_history_count <= 0: - continue - - orig_row = int(row_indices[local_pos].detach().cpu()) - meta_chunk.append( - { - "patient_id": row_base["patient_id"], - "dataset_index": row_base["dataset_index"], - "eid": row_base["eid"], - "sex": row_base["sex"], - "landmark_age": row_base["landmark_age"], - "tau": row_base["tau"], - "followup_end_time": row_base["followup_end_time"], - "history_disease_count": row_base["history_disease_count"], - "selected_disease_history_count": disease_history_count, - "selected_disease_token_id": int(disease_token), - "selected_disease_code": str(disease_meta.get("code", "")), - "selected_disease_name": str(disease_meta.get("name", "")), - "selected_disease_organ_system": str( - disease_meta.get("organ_system", "") - ), - "selected_disease_organ_system_label": str( - disease_meta.get("organ_system_label", "") - ), - "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()), - } + meta_chunk: list[dict[str, Any]] = [] + row_ids = batch_dev["row_idx"][local_rows].detach().cpu().numpy().astype(np.int64) + disease_tokens_list = disease_token_ids.detach().cpu().numpy().astype(np.int64) + for local_pos, (row_idx, disease_token) in enumerate( + zip(row_ids.tolist(), disease_tokens_list.tolist()) + ): + disease_token = int(disease_token) + disease_meta = metadata[disease_token] + row_base = get_row_base(int(row_idx)) + group_counts = row_base["_group_counts"] + disease_history_count = int(row_base["_token_counts"].get(disease_token, 0)) + if disease_history_count <= 0: + raise RuntimeError( + "Internal mismatch: occurred mask selected disease " + f"{disease_token} for row {row_idx}, but cached history has count 0" ) - if meta_chunk: - pending_batch_chunks.append(ablated_chunk) - pending_meta_chunks.append(meta_chunk) - pending_n += len(meta_chunk) - row_offset = row_stop + orig_row = int(local_rows[local_pos].detach().cpu().item()) + meta_chunk.append( + { + "patient_id": row_base["patient_id"], + "dataset_index": row_base["dataset_index"], + "eid": row_base["eid"], + "sex": row_base["sex"], + "landmark_age": row_base["landmark_age"], + "tau": row_base["tau"], + "followup_end_time": row_base["followup_end_time"], + "history_disease_count": row_base["history_disease_count"], + "selected_disease_history_count": disease_history_count, + "selected_disease_token_id": int(disease_token), + "selected_disease_code": str(disease_meta.get("code", "")), + "selected_disease_name": str(disease_meta.get("name", "")), + "selected_disease_organ_system": str( + disease_meta.get("organ_system", "") + ), + "selected_disease_organ_system_label": str( + disease_meta.get("organ_system_label", "") + ), + "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 pending_n >= int(attribution_batch_size): - flush_pending() + if meta_chunk: + pending_batch_chunks.append(ablated_chunk) + pending_meta_chunks.append(meta_chunk) + pending_n += len(meta_chunk) + pair_offset = pair_stop + + if pending_n >= int(attribution_batch_size): + flush_pending() flush_pending() if not shards: