Add assessment and all extra-info experiments
This commit is contained in:
389
train_extra_info_assessment_all_multiseed_linux.sh
Executable file
389
train_extra_info_assessment_all_multiseed_linux.sh
Executable file
@@ -0,0 +1,389 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Train the two remaining extra-information experiments:
|
||||
# 1. assessment_only: 65 routine assessment/body/laboratory variables.
|
||||
# 2. all: all 265 assessment and exposure variables.
|
||||
#
|
||||
# Fixed model:
|
||||
# TrajMixer + all_future + relative + Weibull + timed disease history + sex
|
||||
#
|
||||
# A6000 48 GB defaults:
|
||||
# assessment_only batch_size=256
|
||||
# all batch_size=128
|
||||
#
|
||||
# Each task uses one GPU. Tasks assigned to the same GPU run sequentially;
|
||||
# different GPUs run in parallel.
|
||||
#
|
||||
# Examples:
|
||||
# bash train_extra_info_assessment_all_multiseed_linux.sh --gpus 0
|
||||
# bash train_extra_info_assessment_all_multiseed_linux.sh --gpus 0,1,2
|
||||
# bash train_extra_info_assessment_all_multiseed_linux.sh \
|
||||
# --gpus 0 --seeds 42 --types all --dry-run
|
||||
#
|
||||
|
||||
set -uo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
cd "$SCRIPT_DIR"
|
||||
|
||||
GPU_CSV=""
|
||||
SEED_CSV="42,43,44"
|
||||
TYPE_CSV="assessment_only,all"
|
||||
NUM_WORKERS=4
|
||||
ASSESSMENT_BATCH_SIZE=256
|
||||
ALL_BATCH_SIZE=128
|
||||
PYTHON_BIN="${PYTHON_BIN:-python}"
|
||||
CAMPAIGN_NAME="extra_info_assessment_all_multiseed"
|
||||
DRY_RUN=0
|
||||
|
||||
ENTRYPOINT="$SCRIPT_DIR/train_all_future.py"
|
||||
ASSESSMENT_FILE="$SCRIPT_DIR/extra_info_types_assessment_only.txt"
|
||||
ALL_FILE="$SCRIPT_DIR/extra_info_types_all.txt"
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage:
|
||||
bash train_extra_info_assessment_all_multiseed_linux.sh \
|
||||
--gpus GPU_LIST [options]
|
||||
|
||||
Required:
|
||||
--gpus LIST Comma-separated GPU ids, for example 0,1,2.
|
||||
|
||||
Options:
|
||||
--seeds LIST Comma-separated seeds (default: 42,43,44).
|
||||
--types LIST Subset of assessment_only,all (default: both).
|
||||
--assessment-batch-size N assessment_only batch size (default: 256).
|
||||
--all-batch-size N all batch size (default: 128 for A6000 48 GB).
|
||||
--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 commands without creating files or training.
|
||||
-h, --help Show this help message.
|
||||
|
||||
Fixed experiment settings:
|
||||
architecture traj_mixer_v5
|
||||
target all_future
|
||||
time mode relative
|
||||
distribution weibull
|
||||
disease history timed
|
||||
sex enabled by the model
|
||||
|
||||
A6000 48 GB memory policy:
|
||||
assessment_only batch_size=256
|
||||
all batch_size=128
|
||||
|
||||
If all still runs out of memory because of an unusually long padded batch,
|
||||
restart that experiment with --all-batch-size 64.
|
||||
|
||||
Outputs:
|
||||
runs/<campaign>/seed_<seed>/traj_mixer_v5/...
|
||||
batch_logs/<campaign>/seed_<seed>/<type>.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
|
||||
;;
|
||||
--types)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --types requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
TYPE_CSV="$2"
|
||||
shift 2
|
||||
;;
|
||||
--assessment-batch-size)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --assessment-batch-size requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
ASSESSMENT_BATCH_SIZE="$2"
|
||||
shift 2
|
||||
;;
|
||||
--all-batch-size)
|
||||
[[ $# -ge 2 ]] || {
|
||||
echo "ERROR: --all-batch-size requires a value." >&2
|
||||
exit 2
|
||||
}
|
||||
ALL_BATCH_SIZE="$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
|
||||
}
|
||||
[[ -n "$TYPE_CSV" ]] || {
|
||||
echo "ERROR: --types must not be empty." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$ASSESSMENT_BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || {
|
||||
echo "ERROR: --assessment-batch-size must be a positive integer." >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "$ALL_BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || {
|
||||
echo "ERROR: --all-batch-size must be a positive integer." >&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
|
||||
}
|
||||
for required_file in "$ENTRYPOINT" "$ASSESSMENT_FILE" "$ALL_FILE"; 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
|
||||
|
||||
IFS=',' read -r -a TYPES <<< "$TYPE_CSV"
|
||||
declare -A SEEN_TYPES=()
|
||||
for info_type in "${TYPES[@]}"; do
|
||||
case "$info_type" in
|
||||
assessment_only|all)
|
||||
;;
|
||||
*)
|
||||
echo "ERROR: invalid type: $info_type" >&2
|
||||
echo "Expected assessment_only or all." >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
[[ -z "${SEEN_TYPES[$info_type]+x}" ]] || {
|
||||
echo "ERROR: duplicate type: $info_type" >&2
|
||||
exit 2
|
||||
}
|
||||
SEEN_TYPES["$info_type"]=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_SEEDS=()
|
||||
declare -a JOB_TYPES=()
|
||||
declare -a JOB_BATCH_SIZES=()
|
||||
declare -a JOB_EXTRA_FILES=()
|
||||
|
||||
for seed in "${SEEDS[@]}"; do
|
||||
for info_type in "${TYPES[@]}"; do
|
||||
JOB_SEEDS+=("$seed")
|
||||
JOB_TYPES+=("$info_type")
|
||||
if [[ "$info_type" == "assessment_only" ]]; then
|
||||
JOB_BATCH_SIZES+=("$ASSESSMENT_BATCH_SIZE")
|
||||
JOB_EXTRA_FILES+=("$ASSESSMENT_FILE")
|
||||
else
|
||||
JOB_BATCH_SIZES+=("$ALL_BATCH_SIZE")
|
||||
JOB_EXTRA_FILES+=("$ALL_FILE")
|
||||
fi
|
||||
done
|
||||
done
|
||||
|
||||
print_command() {
|
||||
printf '%q ' "$@"
|
||||
printf '\n'
|
||||
}
|
||||
|
||||
run_job() {
|
||||
local job_index="$1"
|
||||
local gpu="$2"
|
||||
local seed="${JOB_SEEDS[$job_index]}"
|
||||
local info_type="${JOB_TYPES[$job_index]}"
|
||||
local batch_size="${JOB_BATCH_SIZES[$job_index]}"
|
||||
local extra_file="${JOB_EXTRA_FILES[$job_index]}"
|
||||
local seed_runs_root="$RUNS_ROOT/seed_$seed"
|
||||
local seed_log_root="$LOG_ROOT/seed_$seed"
|
||||
local log_file="$seed_log_root/$info_type.log"
|
||||
local -a command=(
|
||||
"$PYTHON_BIN"
|
||||
-u
|
||||
"$ENTRYPOINT"
|
||||
--runs_root "$seed_runs_root"
|
||||
--seed "$seed"
|
||||
--batch_size "$batch_size"
|
||||
--num_workers "$NUM_WORKERS"
|
||||
--device cuda
|
||||
--model_architecture traj_mixer_v5
|
||||
--time_mode relative
|
||||
--dist_mode weibull
|
||||
--disease_history_mode timed
|
||||
--extra_info_types_file "$extra_file"
|
||||
)
|
||||
|
||||
if ((!DRY_RUN)); then
|
||||
mkdir -p "$seed_runs_root" "$seed_log_root"
|
||||
fi
|
||||
|
||||
echo "[$(date '+%F %T')] START seed=$seed type=$info_type gpu=$gpu batch=$batch_size"
|
||||
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 type=$info_type gpu=$gpu"
|
||||
return 0
|
||||
else
|
||||
local exit_code=$?
|
||||
echo "[$(date '+%F %T')] FAIL seed=$seed type=$info_type gpu=$gpu exit=$exit_code" >&2
|
||||
echo " See: $log_file" >&2
|
||||
if [[ "$info_type" == "all" ]]; then
|
||||
echo " If this is CUDA OOM, retry with --all-batch-size 64." >&2
|
||||
fi
|
||||
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_SEEDS[@]}; job_index += ${#GPU_IDS[@]})); do
|
||||
run_job "$job_index" "$gpu" || failed=1
|
||||
done
|
||||
return "$failed"
|
||||
}
|
||||
|
||||
echo "Campaign: $CAMPAIGN_NAME"
|
||||
echo "Seeds: ${SEEDS[*]}"
|
||||
echo "Extra-info types: ${TYPES[*]}"
|
||||
echo "GPUs: ${GPU_IDS[*]}"
|
||||
echo "assessment_only batch size: $ASSESSMENT_BATCH_SIZE"
|
||||
echo "all batch size: $ALL_BATCH_SIZE"
|
||||
echo "Total tasks: ${#JOB_SEEDS[@]}"
|
||||
echo "Runs root: $RUNS_ROOT"
|
||||
echo "Log root: $LOG_ROOT"
|
||||
if command -v nvidia-smi >/dev/null 2>&1; then
|
||||
echo "Selected GPU inventory:"
|
||||
for gpu in "${GPU_IDS[@]}"; do
|
||||
nvidia-smi \
|
||||
--id="$gpu" \
|
||||
--query-gpu=index,name,memory.total \
|
||||
--format=csv,noheader \
|
||||
2>/dev/null || true
|
||||
done
|
||||
fi
|
||||
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 assessment_only/all training tasks completed successfully."
|
||||
fi
|
||||
Reference in New Issue
Block a user