Train assessment smoking and alcohol model

This commit is contained in:
2026-07-31 13:14:35 +08:00
parent efd18f1be2
commit 89dcf4b362
2 changed files with 129 additions and 131 deletions

View File

@@ -0,0 +1,71 @@
# Assessment/body-measurement variables plus smoking and alcohol (field_type=1 plus types 66-67)
# Generated from field_ids_enriched.csv using prepare_data.py other-info type ordering.
# BMI is already included as assessment type 11 and is not duplicated.
# Format: <extra_info_type_id> # <var_name> | <full_name>
1 # waist_circumference | Waist circumference
2 # hip_circumference | Hip circumference
3 # standing_height | Standing height
4 # fasting_time | Fasting time
5 # pulse_rate | Pulse rate automated reading
6 # dbp | Diastolic blood pressure automated reading
7 # sbp | Systolic blood pressure automated reading
8 # fev1_best | Forced expiratory volume in 1-second (FEV1) Best measure
9 # fvc_best | Forced vital capacity (FVC) Best measure
10 # fev1_fvc_ratio | FEV1/ FVC ratio Z-score
11 # bmi | Body mass index (BMI)
12 # WBC | White blood cell (leukocyte) count
13 # RBC | Red blood cell (erythrocyte) count
14 # hemoglobin | Haemoglobin concentration
15 # hematocrit | Haematocrit percentage
16 # MCV | Mean corpuscular volume
17 # MCH | Mean corpuscular haemoglobin
18 # MCHC | Mean corpuscular haemoglobin concentration
19 # Pc | Platelet count
20 # MPV | Mean platelet (thrombocyte) volume
21 # LymC | Lymphocyte count
22 # MonC | Monocyte count
23 # NeuC | Neutrophill count
24 # EosC | Eosinophill count
25 # BasC | Basophill count
26 # nRBC | Nucleated red blood cell count
27 # RC | Reticulocyte count
28 # MRV | Mean reticulocyte volume
29 # MSCV | Mean sphered cell volume
30 # IRF | Immature reticulocyte fraction
31 # HLSRC | High light scatter reticulocyte count
32 # MicU | Microalbumin in urine
33 # CreaU | Creatinine (enzymatic) in urine
34 # PotU | Potassium in urine
35 # SodU | Sodium in urine
36 # Alb | Albumin
37 # ALP | Alkaline phosphatase
38 # Alanine | Alanine aminotransferase
39 # ApoA | Apolipoprotein A
40 # ApoB | Apolipoprotein B
41 # AA | Aspartate aminotransferase
42 # DBil | Direct bilirubin
43 # Urea | Urea
44 # Calcium | Calcium
45 # Cholesterol | Cholesterol
46 # Creatinine | Creatinine
47 # CRP | C-reactive protein
48 # CystatinC | Cystatin C
49 # GGT | Gamma glutamyltransferase
50 # Glu | Glucose
51 # HbA1c | Glycated haemoglobin (HbA1c)
52 # HDL | HDL cholesterol
53 # IGF1 | IGF-1
54 # LDL | LDL direct
55 # LpA | Lipoprotein A
56 # Oestradiol | Oestradiol
57 # Phosphate | Phosphate
58 # Rheu | Rheumatoid factor
59 # SHBG | SHBG
60 # TotalBil | Total bilirubin
61 # Testosterone | Testosterone
62 # TotalProtein | Total protein
63 # Tri | Triglycerides
64 # Urate | Urate
65 # VitaminD | Vitamin D
66 # smoking | Current tobacco smoking
67 # alcohol | Alcohol intake frequency.

View File

@@ -1,15 +1,18 @@
#!/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.
# Train the assessment + smoking + alcohol extra-information experiment.
#
# Fixed model:
# TrajMixer + all_future + relative + Weibull + timed disease history + sex
#
# A6000 48 GB defaults:
# assessment_only batch_size=256
# all batch_size=128
# Extra information:
# - 65 routine assessment/body/laboratory variables
# - smoking
# - alcohol
# - BMI is already included in the assessment variables
#
# A6000 48 GB default:
# batch_size=256
#
# Each task uses one GPU. Tasks assigned to the same GPU run sequentially;
# different GPUs run in parallel.
@@ -18,7 +21,7 @@
# 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
# --gpus 0 --seeds 42 --dry-run
#
set -uo pipefail
@@ -28,17 +31,14 @@ 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
BATCH_SIZE=256
PYTHON_BIN="${PYTHON_BIN:-python}"
CAMPAIGN_NAME="extra_info_assessment_all_multiseed"
CAMPAIGN_NAME="extra_info_assessment_smoking_alcohol_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"
EXTRA_INFO_TYPES_FILE="$SCRIPT_DIR/extra_info_types_assessment_smoking_alcohol.txt"
usage() {
cat <<'EOF'
@@ -51,9 +51,7 @@ Required:
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).
--batch-size N Batch size per task (default: 256).
--num-workers N DataLoader workers per task (default: 4).
--python PATH Python executable (default: $PYTHON_BIN or python).
--campaign NAME Output campaign name.
@@ -67,17 +65,14 @@ Fixed experiment settings:
distribution weibull
disease history timed
sex enabled by the model
extra information extra_info_types_assessment_smoking_alcohol.txt
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.
A6000 48 GB default:
batch_size 256
Outputs:
runs/<campaign>/seed_<seed>/traj_mixer_v5/...
batch_logs/<campaign>/seed_<seed>/<type>.log
batch_logs/<campaign>/seed_<seed>/assessment_smoking_alcohol.log
EOF
}
@@ -99,28 +94,12 @@ while (($# > 0)); do
SEED_CSV="$2"
shift 2
;;
--types)
--batch-size)
[[ $# -ge 2 ]] || {
echo "ERROR: --types requires a value." >&2
echo "ERROR: --batch-size 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"
BATCH_SIZE="$2"
shift 2
;;
--num-workers)
@@ -172,16 +151,8 @@ done
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
[[ "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || {
echo "ERROR: --batch-size must be a positive integer." >&2
exit 2
}
[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || {
@@ -192,12 +163,14 @@ done
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
[[ -f "$ENTRYPOINT" ]] || {
echo "ERROR: missing training entrypoint: $ENTRYPOINT" >&2
exit 2
}
[[ -f "$EXTRA_INFO_TYPES_FILE" ]] || {
echo "ERROR: missing extra-info file: $EXTRA_INFO_TYPES_FILE" >&2
exit 2
}
done
command -v "$PYTHON_BIN" >/dev/null 2>&1 || {
echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
exit 2
@@ -231,86 +204,44 @@ for seed in "${SEEDS[@]}"; do
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 seed="$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 log_file="$seed_log_root/assessment_smoking_alcohol.log"
local -a command=(
"$PYTHON_BIN"
-u
"$ENTRYPOINT"
--runs_root "$seed_runs_root"
--seed "$seed"
--batch_size "$batch_size"
--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"
--extra_info_types_file "$EXTRA_INFO_TYPES_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 "[$(date '+%F %T')] START seed=$seed gpu=$gpu batch=$BATCH_SIZE"
echo " log=$log_file"
if ((DRY_RUN)); then
printf ' CUDA_VISIBLE_DEVICES=%q ' "$gpu"
@@ -320,15 +251,12 @@ run_job() {
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"
echo "[$(date '+%F %T')] DONE seed=$seed 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 "[$(date '+%F %T')] FAIL seed=$seed 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
}
@@ -336,22 +264,21 @@ run_job() {
worker() {
local slot="$1"
local gpu="${GPU_IDS[$slot]}"
local job_index
local seed_index
local failed=0
for ((job_index = slot; job_index < ${#JOB_SEEDS[@]}; job_index += ${#GPU_IDS[@]})); do
run_job "$job_index" "$gpu" || failed=1
for ((seed_index = slot; seed_index < ${#SEEDS[@]}; seed_index += ${#GPU_IDS[@]})); do
run_job "${SEEDS[$seed_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 "Batch size: $BATCH_SIZE"
echo "Extra-info file: $EXTRA_INFO_TYPES_FILE"
echo "Total tasks: ${#SEEDS[@]}"
echo "Runs root: $RUNS_ROOT"
echo "Log root: $LOG_ROOT"
if command -v nvidia-smi >/dev/null 2>&1; then
@@ -385,5 +312,5 @@ fi
if ((DRY_RUN)); then
echo "Dry run completed successfully."
else
echo "All assessment_only/all training tasks completed successfully."
echo "All assessment + smoking + alcohol tasks completed successfully."
fi