From 059c756686038dcb6ee7492d935ec2acf60332d3 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Mon, 27 Jul 2026 10:38:11 +0800 Subject: [PATCH] Add key-model multiseed training script --- train_key_models_multiseed_linux.sh | 357 ++++++++++++++++++++++++++++ 1 file changed, 357 insertions(+) create mode 100644 train_key_models_multiseed_linux.sh diff --git a/train_key_models_multiseed_linux.sh b/train_key_models_multiseed_linux.sh new file mode 100644 index 0000000..d7dd459 --- /dev/null +++ b/train_key_models_multiseed_linux.sh @@ -0,0 +1,357 @@ +#!/usr/bin/env bash +# +# Re-run the key DeepHealth experiments with additional random seeds. +# +# The selected configurations form a compact evidence chain: +# 1. FFN + next_token + absolute + exponential +# Delphi2M reproduction baseline. +# 2. FFN + all_future + absolute + exponential +# Isolates the target change from next_token to all_future. +# 3. TrajMixer + all_future + absolute + exponential +# Isolates the architecture change from FFN to TrajMixer. +# 4. TrajMixer + all_future + relative + exponential +# Isolates the time-mode change and is the current disease-best candidate. +# 5. TrajMixer + all_future + relative + Weibull +# Unified disease/death candidate. +# 6. FFN + all_future + relative + Weibull +# Current mortality-best control and matched architecture comparison for (5). +# +# Each task uses one GPU. Tasks assigned to the same GPU run sequentially, +# while different GPUs run in parallel. +# +# Examples: +# bash train_key_models_multiseed_linux.sh --gpus 0,1,2,3 +# bash train_key_models_multiseed_linux.sh --gpus 0,1 --seeds 43,44 +# + +set -uo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +GPU_CSV="" +SEED_CSV="43,44,45" +NUM_WORKERS=4 +PYTHON_BIN="${PYTHON_BIN:-python}" +CAMPAIGN_NAME="key_models_multiseed_smoking_alcohol_bmi" +DRY_RUN=0 + +BATCH_SIZE=256 +EXTRA_INFO_TYPES_FILE="$SCRIPT_DIR/extra_info_types_smoking_alcohol_bmi.txt" + +usage() { + cat <<'EOF' +Usage: + bash train_key_models_multiseed_linux.sh --gpus GPU_LIST [options] + +Required: + --gpus LIST Comma-separated GPU ids, for example 0,1,2,3. + +Options: + --seeds LIST Additional seeds (default: 43,44,45). + Use 43,44 for two additional seeds. + --num-workers N DataLoader workers per task (default: 4). + --python PATH Python executable (default: $PYTHON_BIN or python). + --campaign NAME Output campaign name. + --dry-run Print all commands without running them. + -h, --help Show this help message. + +Fixed experiment settings: + batch_size 256 + extra_info_types extra_info_types_smoking_alcohol_bmi.txt + model size Defaults from the training entrypoints + tasks per seed 6 + +Outputs: + runs//seed_//... + batch_logs//seed_/.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 + ;; + --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 +} +[[ -f "$EXTRA_INFO_TYPES_FILE" ]] || { + echo "ERROR: missing extra-info file: $EXTRA_INFO_TYPES_FILE" >&2 + exit 2 +} +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)); then + mkdir -p "$RUNS_ROOT" "$LOG_ROOT" +fi + +declare -a JOB_NAMES=() +declare -a JOB_SEEDS=() +declare -a JOB_ENTRYPOINTS=() +declare -a JOB_ARCHITECTURES=() +declare -a JOB_TIME_MODES=() +declare -a JOB_DIST_MODES=() + +add_job() { + JOB_NAMES+=("$1") + JOB_SEEDS+=("$2") + JOB_ENTRYPOINTS+=("$3") + JOB_ARCHITECTURES+=("$4") + JOB_TIME_MODES+=("$5") + JOB_DIST_MODES+=("$6") +} + +for seed in "${SEEDS[@]}"; do + add_job \ + "ffn_next_token_absolute_exponential" \ + "$seed" \ + "train_next_step.py" \ + "transformer_ffn_v1" \ + "absolute" \ + "exponential" + + add_job \ + "ffn_all_future_absolute_exponential" \ + "$seed" \ + "train_all_future.py" \ + "transformer_ffn_v1" \ + "absolute" \ + "exponential" + + add_job \ + "traj_mixer_all_future_absolute_exponential" \ + "$seed" \ + "train_all_future.py" \ + "traj_mixer_v5" \ + "absolute" \ + "exponential" + + add_job \ + "traj_mixer_all_future_relative_exponential" \ + "$seed" \ + "train_all_future.py" \ + "traj_mixer_v5" \ + "relative" \ + "exponential" + + add_job \ + "traj_mixer_all_future_relative_weibull" \ + "$seed" \ + "train_all_future.py" \ + "traj_mixer_v5" \ + "relative" \ + "weibull" + + add_job \ + "ffn_all_future_relative_weibull" \ + "$seed" \ + "train_all_future.py" \ + "transformer_ffn_v1" \ + "relative" \ + "weibull" +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 entrypoint="${JOB_ENTRYPOINTS[$job_index]}" + local architecture="${JOB_ARCHITECTURES[$job_index]}" + local time_mode="${JOB_TIME_MODES[$job_index]}" + local dist_mode="${JOB_DIST_MODES[$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/$entrypoint" + --runs_root "$seed_runs_root" + --seed "$seed" + --batch_size "$BATCH_SIZE" + --extra_info_types_file "$EXTRA_INFO_TYPES_FILE" + --model_architecture "$architecture" + --num_workers "$NUM_WORKERS" + --device cuda + ) + + if ((!DRY_RUN)); then + mkdir -p "$seed_runs_root" "$seed_log_root" + fi + + if [[ "$entrypoint" == "train_all_future.py" ]]; then + command+=( + --time_mode "$time_mode" + --dist_mode "$dist_mode" + ) + 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 "Additional seeds: ${SEEDS[*]}" +echo "GPUs: ${GPU_IDS[*]}" +echo "Configurations per seed: 6" +echo "Total 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 logs under: $LOG_ROOT" >&2 + exit 1 +fi + +if ((DRY_RUN)); then + echo "Dry run completed successfully." +else + echo "All training tasks completed successfully." +fi