144 lines
4.2 KiB
Bash
Executable File
144 lines
4.2 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
# Linux bash 5.2+ training-only script.
|
|
#
|
|
# Based on the existing runs, the objective/time/death-distribution checks are
|
|
# already covered. The remaining gap for the current proof chain is the
|
|
# extra-info ablation under the final candidate model:
|
|
#
|
|
# all_future + relative time + mixed death/risk head
|
|
#
|
|
# This script only launches those missing training jobs. It intentionally does
|
|
# not call evaluate_*.py and does not add extra random seeds. Set
|
|
# MODEL_ARCHITECTURE=traj_mixer_v5 to run the TrajMixer variant.
|
|
|
|
cd "$(dirname "${BASH_SOURCE[0]}")"
|
|
|
|
PYTHON_BIN="${PYTHON_BIN:-python}"
|
|
DEVICE="${DEVICE:-cuda}"
|
|
NUM_WORKERS="${NUM_WORKERS:-4}"
|
|
PROGRESS_INTERVAL="${PROGRESS_INTERVAL:-20}"
|
|
MODEL_ARCHITECTURE="${MODEL_ARCHITECTURE:-transformer_ffn_v1}"
|
|
N_LAYER="${N_LAYER:-12}"
|
|
|
|
TIME_MODE="relative"
|
|
DIST_MODE="mixed"
|
|
SEED="42"
|
|
VALIDATION_QUERY_SEED="42"
|
|
|
|
COMMON_ARGS=(
|
|
--data_prefix ukb
|
|
--labels_file labels.csv
|
|
--seed "${SEED}"
|
|
--validation_query_seed "${VALIDATION_QUERY_SEED}"
|
|
--train_eid_file ukb_train_eid.csv
|
|
--val_eid_file ukb_val_eid.csv
|
|
--test_eid_file ukb_test_eid.csv
|
|
--min_history_events 1
|
|
--min_future_events 1
|
|
--n_embd 120
|
|
--n_head 10
|
|
--n_layer "${N_LAYER}"
|
|
--model_architecture "${MODEL_ARCHITECTURE}"
|
|
--n_bins 16
|
|
--extra_pool_reduce mean
|
|
--dropout 0.0
|
|
--batch_size 256
|
|
--base_lr 0.0003
|
|
--weight_decay 0.1
|
|
--betas 0.9 0.99
|
|
--grad_clip 1.0
|
|
--max_epochs 200
|
|
--warmup_epochs 10
|
|
--patience 15
|
|
--min_lr_ratio 0.1
|
|
--num_workers "${NUM_WORKERS}"
|
|
--device "${DEVICE}"
|
|
--progress_interval "${PROGRESS_INTERVAL}"
|
|
)
|
|
|
|
already_trained() {
|
|
local extra_file="$1"
|
|
"${PYTHON_BIN}" - "$TIME_MODE" "$DIST_MODE" "$extra_file" "$SEED" "$VALIDATION_QUERY_SEED" "$MODEL_ARCHITECTURE" "$N_LAYER" <<'PY'
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
(
|
|
time_mode,
|
|
dist_mode,
|
|
extra_file,
|
|
seed,
|
|
validation_query_seed,
|
|
model_architecture,
|
|
n_layer,
|
|
) = sys.argv[1:8]
|
|
extra_name = Path(extra_file).name
|
|
|
|
for config_path in Path("runs").rglob("train_config.json"):
|
|
try:
|
|
cfg = json.loads(config_path.read_text(encoding="utf-8"))
|
|
except Exception:
|
|
continue
|
|
|
|
observed_query_seed = cfg.get(
|
|
"all_future_validation_query_seed",
|
|
cfg.get("validation_query_seed", -1),
|
|
)
|
|
|
|
if (
|
|
cfg.get("model_target_mode") == "all_future"
|
|
and cfg.get("model_architecture") == model_architecture
|
|
and int(cfg.get("n_layer", -1)) == int(n_layer)
|
|
and cfg.get("time_mode") == time_mode
|
|
and cfg.get("dist_mode") == dist_mode
|
|
and Path(str(cfg.get("extra_info_types_file", ""))).name == extra_name
|
|
and int(cfg.get("seed", -1)) == int(seed)
|
|
and int(observed_query_seed) == int(validation_query_seed)
|
|
):
|
|
print(config_path.parent)
|
|
raise SystemExit(0)
|
|
|
|
raise SystemExit(1)
|
|
PY
|
|
}
|
|
|
|
train_if_missing() {
|
|
local label="$1"
|
|
local extra_file="$2"
|
|
|
|
if [[ ! -f "${extra_file}" ]]; then
|
|
echo "ERROR: missing extra-info type file: ${extra_file}" >&2
|
|
return 2
|
|
fi
|
|
|
|
echo "==> Checking ${label}: ${MODEL_ARCHITECTURE} n_layer=${N_LAYER} ${TIME_MODE} ${DIST_MODE} all_future with ${extra_file}"
|
|
if existing_run="$(already_trained "$extra_file")"; then
|
|
echo " skip: already trained at ${existing_run}"
|
|
return 0
|
|
fi
|
|
|
|
echo " train: ${label}"
|
|
"${PYTHON_BIN}" train_all_future.py \
|
|
"${COMMON_ARGS[@]}" \
|
|
--time_mode "${TIME_MODE}" \
|
|
--dist_mode "${DIST_MODE}" \
|
|
--extra_info_types_file "${extra_file}"
|
|
}
|
|
|
|
# Already present in runs/:
|
|
# - next-token objective checks under SAB, plus older absolute extra ablations
|
|
# - all-future absolute/relative x exponential/weibull/mixed under SAB
|
|
#
|
|
# Still needed:
|
|
# - final all-future relative+mixed extra-info ablations beyond the existing
|
|
# SAB baseline. These close the disease-only question without expanding seed
|
|
# count or running downstream evaluation.
|
|
train_if_missing "true_disease_only" "extra_info_types_none.txt"
|
|
train_if_missing "assessment_only_extra" "extra_info_types_assessment_only.txt"
|
|
train_if_missing "exposure_only_extra" "extra_info_types_exposure_only.txt"
|
|
train_if_missing "all_extra_info" "extra_info_types_all.txt"
|
|
|
|
echo "All requested training-only missing configurations are done."
|