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

Parallelize extra-info summary reduction

Browse Source
This commit is contained in:
Jiarui Li
2026-07-01 13:21:08 +08:00
parent e3c9ecd19f
commit f417a91a74
2 changed files with 128 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

133
evaluate_extra_info_attribution.py
View File

@@ -1,4 +1,4 @@
"""Evaluate extra-info attribution to death parameters and future disease risks.
"""Evaluate extra-info attribution to death and disease distribution parameters.
For each landmark query, this script scans selected extra-info types that are
available at or before the query age. For each such type it re-runs the model
@@ -15,6 +15,7 @@ from __future__ import annotations
import argparse
import json
import re
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Sequence
@@ -462,6 +463,98 @@ def update_disease_parameter_summary_from_group_stats(
acc[f"sumsq__{column}"] += float(sumsq[int(row_idx), int(group_idx), int(col_idx)])
def merge_summary_dict(
dst: dict[tuple[Any, ...], dict[str, float]],
src: dict[tuple[Any, ...], dict[str, float]],
) -> None:
for key, src_acc in src.items():
dst_acc = dst.setdefault(key, {name: 0.0 for name in src_acc})
for name, value in src_acc.items():
dst_acc[name] = dst_acc.get(name, 0.0) + float(value)
def reduce_attribution_chunk_bundle(
payload: tuple[
list[tuple[pd.DataFrame, np.ndarray]],
list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]],
list[str],
list[str],
],
) -> tuple[dict[tuple[Any, ...], dict[str, float]], dict[tuple[Any, ...], dict[str, float]]]:
death_items, disease_items, group_names, group_labels = payload
death_summary: dict[tuple[Any, ...], dict[str, float]] = {}
disease_summary: dict[tuple[Any, ...], dict[str, float]] = {}
for key_rows, values in death_items:
update_death_summary(
death_summary,
key_rows=key_rows,
values=values,
)
for key_rows, sums, sumsq, counts in disease_items:
update_disease_parameter_summary_from_group_stats(
disease_summary,
key_rows=key_rows,
group_names=group_names,
group_labels=group_labels,
sums=sums,
sumsq=sumsq,
counts=counts,
)
return death_summary, disease_summary
def reduce_attribution_chunks(
*,
death_key_chunks: list[pd.DataFrame],
death_value_chunks: list[np.ndarray],
disease_stat_chunks: list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]],
group_names: list[str],
group_labels: list[str],
cpu_reduce_workers: int,
) -> tuple[dict[tuple[Any, ...], dict[str, float]], dict[tuple[Any, ...], dict[str, float]]]:
n_chunks = max(len(death_key_chunks), len(disease_stat_chunks))
if n_chunks == 0:
return {}, {}
worker_count = max(1, min(int(cpu_reduce_workers), n_chunks))
if worker_count == 1:
return reduce_attribution_chunk_bundle(
(
list(zip(death_key_chunks, death_value_chunks)),
disease_stat_chunks,
group_names,
group_labels,
)
)
bundles = []
for worker_idx in range(worker_count):
start = worker_idx * n_chunks // worker_count
stop = (worker_idx + 1) * n_chunks // worker_count
if start >= stop:
continue
death_items = [
(death_key_chunks[i], death_value_chunks[i])
for i in range(start, min(stop, len(death_key_chunks)))
]
disease_items = disease_stat_chunks[start:min(stop, len(disease_stat_chunks))]
bundles.append((death_items, disease_items, group_names, group_labels))
merged_death: dict[tuple[Any, ...], dict[str, float]] = {}
merged_disease: dict[tuple[Any, ...], dict[str, float]] = {}
with ProcessPoolExecutor(max_workers=len(bundles)) as executor:
futures = [executor.submit(reduce_attribution_chunk_bundle, bundle) for bundle in bundles]
for future in tqdm(as_completed(futures), total=len(futures), desc="CPU summary reduction", dynamic_ncols=True):
death_part, disease_part = future.result()
merge_summary_dict(merged_death, death_part)
merge_summary_dict(merged_disease, disease_part)
return merged_death, merged_disease
def write_death_summary_csv(
path: Path,
summary: dict[tuple[Any, ...], dict[str, float]],
@@ -532,7 +625,7 @@ def write_disease_parameter_summary_csv(
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Compute extra-info ablation attribution for death parameters and future disease risks."
description="Compute extra-info ablation attribution for death and disease distribution parameters."
)
parser.add_argument("--run_path", type=str, required=True)
parser.add_argument(
@@ -563,6 +656,12 @@ def parse_args() -> argparse.Namespace:
help="Forward batch size for expanded extra-info ablation queries.",
)
parser.add_argument("--num_workers", type=int, default=None)
parser.add_argument(
"--cpu_reduce_workers",
type=int,
default=None,
help="Worker processes for CPU-side summary reduction. Defaults to --num_workers.",
)
parser.add_argument("--device", type=str, default=None)
return parser.parse_args()
@@ -685,6 +784,13 @@ def main() -> None:
raise ValueError("attribution_batch_size must be positive")
num_workers = int(cfg_get(args, cfg, "num_workers", 4))
cpu_reduce_workers = int(
args.cpu_reduce_workers
if args.cpu_reduce_workers is not None
else max(1, num_workers)
)
if cpu_reduce_workers <= 0:
raise ValueError("--cpu_reduce_workers must be positive")
loader = DataLoader(
IndexedLandmarkDataset(landmark_dataset),
batch_size=batch_size,
@@ -713,10 +819,9 @@ def main() -> None:
print(f"Extra-info types: {selected_extra_info_types}")
print(f"Landmark rows: {len(landmark_dataset)}")
print(f"Attribution batch size: {attribution_batch_size}")
print(f"CPU reduce workers: {cpu_reduce_workers}")
print(f"Output directory: {output_dir}")
death_summary: dict[tuple[Any, ...], dict[str, float]] = {}
disease_parameter_summary: dict[tuple[Any, ...], dict[str, float]] = {}
death_key_chunks: list[pd.DataFrame] = []
death_value_chunks: list[np.ndarray] = []
disease_stat_chunks: list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]] = []
@@ -799,23 +904,13 @@ def main() -> None:
)
disease_stat_chunks.append((key_table, sums, sumsq, counts))
death_summary.clear()
for key_rows, values in zip(death_key_chunks, death_value_chunks):
update_death_summary(
death_summary,
key_rows=key_rows,
values=values,
)
for key_rows, sums, sumsq, counts in disease_stat_chunks:
update_disease_parameter_summary_from_group_stats(
disease_parameter_summary,
key_rows=key_rows,
death_summary, disease_parameter_summary = reduce_attribution_chunks(
death_key_chunks=death_key_chunks,
death_value_chunks=death_value_chunks,
disease_stat_chunks=disease_stat_chunks,
group_names=group_names,
group_labels=group_labels,
sums=sums,
sumsq=sumsq,
counts=counts,
cpu_reduce_workers=cpu_reduce_workers,
)
death_summary_path = output_dir / "summary_extra_info_death_parameters.csv"

12
run_missing_evaluations.sh
View File

@@ -10,6 +10,7 @@ PYTHON_BIN="${PYTHON_BIN:-python}"
DEVICE="${DEVICE:-cuda}"
EVAL_SPLIT="${EVAL_SPLIT:-test}"
NUM_WORKERS="${NUM_WORKERS:-4}"
CPU_REDUCE_WORKERS="${CPU_REDUCE_WORKERS:-}"
NUM_WORKERS_AUC="${NUM_WORKERS_AUC:-}"
BATCH_SIZE="${BATCH_SIZE:-}"
DATASET_SUBSET_SIZE="${DATASET_SUBSET_SIZE:-}"
@@ -41,6 +42,12 @@ auc_args() {
fi
}
cpu_reduce_args() {
if [[ -n "${CPU_REDUCE_WORKERS}" ]]; then
printf '%s\n' --cpu_reduce_workers "${CPU_REDUCE_WORKERS}"
fi
}
has_completed_dir() {
local dir="$1"
shift
@@ -127,6 +134,9 @@ for run_path in runs/*; do
auc_extra=()
while IFS= read -r arg; do auc_extra+=("${arg}"); done < <(auc_args)
cpu_reduce_extra=()
while IFS= read -r arg; do cpu_reduce_extra+=("${arg}"); done < <(cpu_reduce_args)
run_dir_result_if_missing \
"evaluate_auc.py" \
"${run_path}" \
@@ -153,7 +163,7 @@ for run_path in runs/*; do
"${run_path}/extra_info_attribution_${EVAL_SPLIT}" \
"manifest.json" \
"summary_extra_info_disease_parameters.csv" \
"${PYTHON_BIN}" evaluate_extra_info_attribution.py "${common[@]}"
"${PYTHON_BIN}" evaluate_extra_info_attribution.py "${common[@]}" "${cpu_reduce_extra[@]}"
else
echo " skip evaluate_extra_info_attribution.py: run has no extra-info types"
fi