Files
DeepHealth/run_all_future_first_onset_experiments_linux.sh

391 lines
10 KiB
Bash
Raw Normal View History

#!/usr/bin/env bash
#
# Run the matched experiments required after the all-future first-onset update.
#
# Shared by every experiment:
# - corrected first-onset likelihood and outcome-specific risk exposure;
# - patient/interval/time-uniform query sampling in train/valid/test;
# - timed disease history and smoking/alcohol/BMI extra information.
#
# Per seed, the script trains:
# 1. TrajMixer + relative Weibull + fitted risk baseline (primary model);
# 2. the primary model without the fitted risk baseline (baseline ablation);
# 3. TrajMixer + relative exponential (time-distribution control);
# 4. FFN + relative Weibull (architecture control).
#
# After training, AUC and calibration/point-process-NLL evaluation run
# automatically unless --train-only is supplied.
set -uo pipefail
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"
GPU_CSV=""
SEED_CSV="42,43,44"
NUM_WORKERS=4
PYTHON_BIN="${PYTHON_BIN:-python}"
CAMPAIGN_NAME="all_future_first_onset_v2_multiseed"
DRY_RUN=0
TRAIN_ONLY=0
BATCH_SIZE=256
EXTRA_INFO_TYPES_FILE="$SCRIPT_DIR/extra_info_types_smoking_alcohol_bmi.txt"
TRAIN_EID_FILE="$SCRIPT_DIR/ukb_train_eid.csv"
VAL_EID_FILE="$SCRIPT_DIR/ukb_val_eid.csv"
TEST_EID_FILE="$SCRIPT_DIR/ukb_test_eid.csv"
usage() {
cat <<'EOF'
Usage:
bash run_all_future_first_onset_experiments_linux.sh --gpus LIST [options]
Required:
--gpus LIST Comma-separated GPU ids, for example 0 or 0,1,2,3.
Options:
--seeds LIST Comma-separated seeds (default: 42,43,44).
--num-workers N DataLoader workers per active GPU (default: 4).
--python PATH Python executable (default: $PYTHON_BIN or python).
--campaign NAME Output campaign directory name.
--train-only Skip AUC and calibration/NLL evaluation.
--dry-run Print commands without running them.
-h, --help Show this help message.
Fixed settings:
batch_size 256
extra information smoking + alcohol + BMI
patient split fixed train/validation/test EID files
disease history timed
experiments per seed 4
Outputs:
runs/<campaign>/seed_<seed>/<architecture>/...
batch_logs/<campaign>/seed_<seed>/<experiment>.log
EOF
}
while (($# > 0)); do
case "$1" in
--gpus)
[[ $# -ge 2 ]] || {
echo "ERROR: --gpus requires a value." >&2
exit 2
}
GPU_CSV="$2"
shift 2
;;
--seeds)
[[ $# -ge 2 ]] || {
echo "ERROR: --seeds requires a value." >&2
exit 2
}
SEED_CSV="$2"
shift 2
;;
--num-workers)
[[ $# -ge 2 ]] || {
echo "ERROR: --num-workers requires a value." >&2
exit 2
}
NUM_WORKERS="$2"
shift 2
;;
--python)
[[ $# -ge 2 ]] || {
echo "ERROR: --python requires a value." >&2
exit 2
}
PYTHON_BIN="$2"
shift 2
;;
--campaign)
[[ $# -ge 2 ]] || {
echo "ERROR: --campaign requires a value." >&2
exit 2
}
CAMPAIGN_NAME="$2"
shift 2
;;
--train-only)
TRAIN_ONLY=1
shift
;;
--dry-run)
DRY_RUN=1
shift
;;
-h|--help)
usage
exit 0
;;
*)
echo "ERROR: unknown argument: $1" >&2
usage >&2
exit 2
;;
esac
done
[[ -n "$GPU_CSV" ]] || {
echo "ERROR: --gpus is required." >&2
usage >&2
exit 2
}
[[ -n "$SEED_CSV" ]] || {
echo "ERROR: --seeds must not be empty." >&2
exit 2
}
[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || {
echo "ERROR: --num-workers must be a non-negative integer." >&2
exit 2
}
[[ "$CAMPAIGN_NAME" =~ ^[A-Za-z0-9._-]+$ ]] || {
echo "ERROR: --campaign may contain only letters, numbers, ., _, and -." >&2
exit 2
}
required_files=(
"$SCRIPT_DIR/train_all_future.py"
"$SCRIPT_DIR/evaluate_all_runs_linux.sh"
"$SCRIPT_DIR/evaluate_calibration_all_runs_linux.sh"
"$EXTRA_INFO_TYPES_FILE"
"$TRAIN_EID_FILE"
"$VAL_EID_FILE"
"$TEST_EID_FILE"
)
for required_file in "${required_files[@]}"; do
[[ -f "$required_file" ]] || {
echo "ERROR: missing required file: $required_file" >&2
exit 2
}
done
command -v "$PYTHON_BIN" >/dev/null 2>&1 || {
echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
exit 2
}
IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
declare -A SEEN_GPUS=()
for gpu in "${GPU_IDS[@]}"; do
[[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || {
echo "ERROR: invalid GPU id: $gpu" >&2
exit 2
}
[[ -z "${SEEN_GPUS[$gpu]+x}" ]] || {
echo "ERROR: duplicate GPU id: $gpu" >&2
exit 2
}
SEEN_GPUS["$gpu"]=1
done
IFS=',' read -r -a SEEDS <<< "$SEED_CSV"
declare -A SEEN_SEEDS=()
for seed in "${SEEDS[@]}"; do
[[ "$seed" =~ ^[0-9]+$ ]] || {
echo "ERROR: invalid seed: $seed" >&2
exit 2
}
[[ -z "${SEEN_SEEDS[$seed]+x}" ]] || {
echo "ERROR: duplicate seed: $seed" >&2
exit 2
}
SEEN_SEEDS["$seed"]=1
done
RUNS_ROOT="$SCRIPT_DIR/runs/$CAMPAIGN_NAME"
LOG_ROOT="$SCRIPT_DIR/batch_logs/$CAMPAIGN_NAME"
if ((DRY_RUN == 0)); then
mkdir -p "$RUNS_ROOT" "$LOG_ROOT"
fi
declare -a JOB_NAMES=()
declare -a JOB_SEEDS=()
declare -a JOB_ARCHITECTURES=()
declare -a JOB_DIST_MODES=()
declare -a JOB_BASELINES=()
add_job() {
JOB_NAMES+=("$1")
JOB_SEEDS+=("$2")
JOB_ARCHITECTURES+=("$3")
JOB_DIST_MODES+=("$4")
JOB_BASELINES+=("$5")
}
for seed in "${SEEDS[@]}"; do
add_job \
"traj_mixer_relative_weibull_first_onset" \
"$seed" \
"traj_mixer_v5" \
"weibull" \
"on"
add_job \
"traj_mixer_relative_weibull_no_rate_baseline" \
"$seed" \
"traj_mixer_v5" \
"weibull" \
"off"
add_job \
"traj_mixer_relative_exponential_first_onset" \
"$seed" \
"traj_mixer_v5" \
"exponential" \
"on"
add_job \
"ffn_relative_weibull_first_onset" \
"$seed" \
"transformer_ffn_v1" \
"weibull" \
"on"
done
print_command() {
printf '%q ' "$@"
printf '\n'
}
run_job() {
local job_index="$1"
local gpu="$2"
local job_name="${JOB_NAMES[$job_index]}"
local seed="${JOB_SEEDS[$job_index]}"
local architecture="${JOB_ARCHITECTURES[$job_index]}"
local dist_mode="${JOB_DIST_MODES[$job_index]}"
local baseline="${JOB_BASELINES[$job_index]}"
local seed_runs_root="$RUNS_ROOT/seed_$seed"
local seed_log_root="$LOG_ROOT/seed_$seed"
local log_file="$seed_log_root/$job_name.log"
local -a command=(
"$PYTHON_BIN"
-u
"$SCRIPT_DIR/train_all_future.py"
--runs_root "$seed_runs_root"
--seed "$seed"
--batch_size "$BATCH_SIZE"
--num_workers "$NUM_WORKERS"
--device cuda
--model_architecture "$architecture"
--time_mode relative
--dist_mode "$dist_mode"
--disease_history_mode timed
--extra_info_types_file "$EXTRA_INFO_TYPES_FILE"
--train_eid_file "$TRAIN_EID_FILE"
--val_eid_file "$VAL_EID_FILE"
--test_eid_file "$TEST_EID_FILE"
)
if [[ "$baseline" == "on" ]]; then
command+=(--risk_head_bias)
else
command+=(--no-risk_head_bias)
fi
if ((DRY_RUN == 0)); then
mkdir -p "$seed_runs_root" "$seed_log_root"
fi
echo "[$(date '+%F %T')] START seed=$seed job=$job_name gpu=$gpu"
echo " log=$log_file"
if ((DRY_RUN)); then
printf ' CUDA_VISIBLE_DEVICES=%q ' "$gpu"
print_command "${command[@]}"
return 0
fi
if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \
"${command[@]}" >"$log_file" 2>&1; then
echo "[$(date '+%F %T')] DONE seed=$seed job=$job_name gpu=$gpu"
return 0
else
local exit_code=$?
echo "[$(date '+%F %T')] FAIL seed=$seed job=$job_name gpu=$gpu exit=$exit_code" >&2
echo " See: $log_file" >&2
return "$exit_code"
fi
}
worker() {
local slot="$1"
local gpu="${GPU_IDS[$slot]}"
local job_index
local failed=0
for ((job_index = slot; job_index < ${#JOB_NAMES[@]}; job_index += ${#GPU_IDS[@]})); do
run_job "$job_index" "$gpu" || failed=1
done
return "$failed"
}
echo "Campaign: $CAMPAIGN_NAME"
echo "Seeds: ${SEEDS[*]}"
echo "GPUs: ${GPU_IDS[*]}"
echo "Experiments per seed: 4"
echo "Total training tasks: ${#JOB_NAMES[@]}"
echo "Runs root: $RUNS_ROOT"
echo "Log root: $LOG_ROOT"
echo
declare -a WORKER_PIDS=()
for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do
worker "$slot" &
WORKER_PIDS+=("$!")
done
overall_status=0
for pid in "${WORKER_PIDS[@]}"; do
wait "$pid" || overall_status=1
done
if ((overall_status != 0)); then
echo "One or more training tasks failed. Inspect: $LOG_ROOT" >&2
exit 1
fi
if ((TRAIN_ONLY == 0)); then
evaluation_status=0
auc_command=(
bash
"$SCRIPT_DIR/evaluate_all_runs_linux.sh"
--gpus "$GPU_CSV"
--runs-root "$RUNS_ROOT"
--log-root "$LOG_ROOT/evaluate_auc"
--python "$PYTHON_BIN"
--num-workers "$NUM_WORKERS"
--num-workers-auc "$NUM_WORKERS"
)
calibration_command=(
bash
"$SCRIPT_DIR/evaluate_calibration_all_runs_linux.sh"
--gpus "$GPU_CSV"
--runs-root "$RUNS_ROOT"
--log-root "$LOG_ROOT/evaluate_calibration"
--python "$PYTHON_BIN"
--num-workers "$NUM_WORKERS"
--num-workers-calibration "$NUM_WORKERS"
)
echo ">> AUC evaluation"
print_command "${auc_command[@]}"
if ((DRY_RUN == 0)); then
"${auc_command[@]}" || evaluation_status=1
fi
echo ">> Calibration and point-process NLL evaluation"
print_command "${calibration_command[@]}"
if ((DRY_RUN == 0)); then
"${calibration_command[@]}" || evaluation_status=1
fi
if ((evaluation_status != 0)); then
echo "One or more evaluation workflows failed. Inspect: $LOG_ROOT" >&2
exit 1
fi
fi
if ((DRY_RUN)); then
echo "Dry run completed successfully."
else
echo "All required all-future experiments completed successfully."
fi