Vectorize per-disease attribution batches
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user