From 2ae92da0eca359fb916da02706658efd0800d799 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Sat, 27 Jun 2026 14:08:02 +0800 Subject: [PATCH] Stream mortality attribution batches --- evaluate_event_free_survival.py | 233 ++++++++++++++++---------------- 1 file changed, 120 insertions(+), 113 deletions(-) diff --git a/evaluate_event_free_survival.py b/evaluate_event_free_survival.py index 64b2043..57497ed 100644 --- a/evaluate_event_free_survival.py +++ b/evaluate_event_free_survival.py @@ -288,82 +288,57 @@ def collate_indexed_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[st } -def build_group_ablated_batch( +def build_group_ablated_slice( batch: Dict[str, torch.Tensor], - group_names: Sequence[str], - organ_groups: dict[str, list[int]], - occurred: torch.Tensor, -) -> tuple[Dict[str, torch.Tensor], list[str]]: - """Build a group-major batch for organ/system mortality attribution.""" - active_groups: list[str] = [] - for group in group_names: - ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=occurred.device) - if ids.numel() > 0 and bool(occurred[:, ids].any().item()): - active_groups.append(group) - - if not active_groups: - return {}, [] - - event_rows: list[torch.Tensor] = [] - time_rows: list[torch.Tensor] = [] - readout_rows: list[torch.Tensor] = [] - landmark_positions: list[torch.Tensor] = [] - + token_ids: Sequence[int], + row_start: int, + row_stop: int, +) -> Dict[str, torch.Tensor]: + """Build one fixed-width ablated slice without rebuilding variable-length rows.""" event_seq = batch["event_seq"] - time_seq = batch["time_seq"] - readout_mask = batch["readout_mask"] - padding_mask = batch["padding_mask"].bool() - for group in active_groups: - selected = {int(token) for token in organ_groups[group]} - for i in range(event_seq.shape[0]): - valid = padding_mask[i] - events = event_seq[i, valid] - times = time_seq[i, valid] - reads = readout_mask[i, valid] - keep = torch.ones_like(events, dtype=torch.bool) - for token in selected: - keep &= events != int(token) - kept_events = events[keep] - kept_times = times[keep] - kept_reads = reads[keep] - if kept_events.numel() == 0: - kept_events = torch.tensor( - [CHECKUP_IDX], - dtype=event_seq.dtype, - device=event_seq.device, - ) - kept_times = batch["t_query"][i : i + 1].to( - dtype=time_seq.dtype, - device=time_seq.device, - ) - kept_reads = torch.ones(1, dtype=torch.bool, device=readout_mask.device) - - if bool(kept_reads.any()): - landmark_pos = torch.nonzero(kept_reads, as_tuple=False)[-1, 0] - else: - landmark_pos = torch.tensor( - int(kept_events.numel() - 1), - dtype=batch["landmark_pos"].dtype, - device=batch["landmark_pos"].device, - ) - kept_reads = torch.zeros_like(kept_events, dtype=torch.bool) - kept_reads[int(landmark_pos.item())] = True - - event_rows.append(kept_events) - time_rows.append(kept_times) - readout_rows.append(kept_reads) - landmark_positions.append(landmark_pos.to(dtype=batch["landmark_pos"].dtype)) - - repeat_count = len(active_groups) out: Dict[str, torch.Tensor] = {} - out["event_seq"] = pad_sequence(event_rows, batch_first=True, padding_value=PAD_IDX) - out["time_seq"] = pad_sequence(time_rows, batch_first=True, padding_value=0.0) - out["readout_mask"] = pad_sequence( - readout_rows, batch_first=True, padding_value=False + out["event_seq"] = event_seq[row_start:row_stop].clone() + out["time_seq"] = batch["time_seq"][row_start:row_stop] + out["readout_mask"] = batch["readout_mask"][row_start:row_stop].clone() + out["padding_mask"] = batch["padding_mask"][row_start:row_stop].bool().clone() + out["landmark_pos"] = batch["landmark_pos"][row_start:row_stop].clone() + + seq_len = int(event_seq.shape[1]) + positions = torch.arange(seq_len, device=event_seq.device)[None, :] + ids = torch.as_tensor(token_ids, dtype=event_seq.dtype, device=event_seq.device) + remove = torch.isin(out["event_seq"], ids) & out["padding_mask"] + out["event_seq"] = torch.where( + remove, + torch.full_like(out["event_seq"], PAD_IDX), + out["event_seq"], ) - out["padding_mask"] = out["event_seq"] > PAD_IDX - out["landmark_pos"] = torch.stack(landmark_positions) + 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_start + 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", @@ -379,18 +354,62 @@ def build_group_ablated_batch( "row_idx", ) for key in repeated_keys: - value = batch[key] - repeats = (repeat_count,) + (1,) * (value.ndim - 1) - out[key] = value.repeat(repeats) - return out, active_groups + out[key] = batch[key][row_start:row_stop] + return out -def slice_tensor_batch( +def concat_tensor_batches(chunks: Sequence[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]: + return { + key: torch.cat([chunk[key] for chunk in chunks], dim=0) + for key in chunks[0] + } + + +def iter_group_ablated_batches( batch: Dict[str, torch.Tensor], - start: int, - stop: int, -) -> Dict[str, torch.Tensor]: - return {key: value[start:stop] for key, value in batch.items()} + group_names: Sequence[str], + organ_groups: dict[str, list[int]], + occurred: torch.Tensor, + max_batch_size: int, +): + """Yield ablated chunks as soon as enough rows are available for a forward pass.""" + batch_n = int(batch["event_seq"].shape[0]) + pending_batches: list[Dict[str, torch.Tensor]] = [] + pending_groups: list[str] = [] + pending_rows: list[int] = [] + pending_n = 0 + + for group in group_names: + ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=occurred.device) + if ids.numel() == 0 or not bool(occurred[:, ids].any().item()): + continue + + row_start = 0 + while row_start < batch_n: + capacity = int(max_batch_size) - pending_n + row_stop = min(batch_n, row_start + capacity) + chunk = build_group_ablated_slice( + batch=batch, + token_ids=organ_groups[group], + row_start=row_start, + row_stop=row_stop, + ) + chunk_n = int(row_stop - row_start) + pending_batches.append(chunk) + pending_groups.extend([group] * chunk_n) + pending_rows.extend(range(row_start, row_stop)) + pending_n += chunk_n + row_start = row_stop + + if pending_n >= int(max_batch_size): + yield concat_tensor_batches(pending_batches), pending_groups, pending_rows + pending_batches = [] + pending_groups = [] + pending_rows = [] + pending_n = 0 + + if pending_batches: + yield concat_tensor_batches(pending_batches), pending_groups, pending_rows @torch.no_grad() @@ -646,6 +665,7 @@ def main() -> None: print(f"Landmark ages: {landmark_ages.tolist()}") print(f"Tau: {tau:g} years") print(f"Dist mode: {dist_mode}") + print(f"Device: {device}") print(f"Death token: {death_idx}") print(f"Organ/system groups: {len(group_names)}") print(f"Landmark rows: {len(landmark_dataset)}") @@ -698,47 +718,34 @@ def main() -> None: group_mortality_attr_prob[group] = zeros.copy() group_mortality_attr_hazard[group] = zeros.copy() - ablated_batch, active_groups = build_group_ablated_batch( + for ablated_chunk, chunk_groups, chunk_rows in iter_group_ablated_batches( batch=batch, group_names=group_names, organ_groups=organ_groups, occurred=occurred, - ) - if active_groups: - ablated_risk_chunks: list[torch.Tensor] = [] - expanded_n = int(ablated_batch["event_seq"].shape[0]) - for start in range(0, expanded_n, attribution_batch_size): - chunk = slice_tensor_batch( - ablated_batch, - start, - min(start + attribution_batch_size, expanded_n), - ) - ablated_risk_chunks.append( - death_risk_for_batch( - model=model, - batch=chunk, - device=device, - model_target_mode=model_target_mode, - readout_name=readout_name, - readout_reduce=readout_reduce, - dist_mode=dist_mode, - tau=tau, - ) - ) - - ablated_death_risk = torch.cat(ablated_risk_chunks, dim=0).view( - len(active_groups), batch_n + max_batch_size=attribution_batch_size, + ): + ablated_death_risk = death_risk_for_batch( + model=model, + batch=ablated_chunk, + device=device, + model_target_mode=model_target_mode, + readout_name=readout_name, + readout_reduce=readout_reduce, + dist_mode=dist_mode, + tau=tau, ) + row_tensor = torch.as_tensor(chunk_rows, dtype=torch.long, device=device) ablated_death_hazard = mortality_hazard_from_risk(ablated_death_risk) - active_attr_hazard = ( - death_hazard_tensor[None, :] - ablated_death_hazard + attr_prob = ( + death_risk_tensor[row_tensor] - ablated_death_risk ).detach().cpu().numpy() - active_attr_prob = ( - death_risk_tensor[None, :] - ablated_death_risk + attr_hazard = ( + death_hazard_tensor[row_tensor] - ablated_death_hazard ).detach().cpu().numpy() - for group_idx, group in enumerate(active_groups): - group_mortality_attr_prob[group] = active_attr_prob[group_idx] - group_mortality_attr_hazard[group] = active_attr_hazard[group_idx] + for local_idx, (group, row_idx) in enumerate(zip(chunk_groups, chunk_rows)): + group_mortality_attr_prob[group][row_idx] = attr_prob[local_idx] + group_mortality_attr_hazard[group][row_idx] = attr_hazard[local_idx] row_indices = batch["row_idx"].cpu().numpy().astype(np.int64) for j, row_idx in enumerate(row_indices.tolist()):