89 lines
2.5 KiB
Bash
Executable File
89 lines
2.5 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
|
|
set -euo pipefail
|
|
|
|
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
|
DATA_DIR="${ROOT_DIR}/data"
|
|
OUTPUT_FILE="${ROOT_DIR}/sleep-model.pt"
|
|
MODEL_FILE="${OUTPUT_FILE}"
|
|
|
|
if [[ ${1:-} == "-h" || ${1:-} == "--help" ]]; then
|
|
cat <<'EOF'
|
|
Usage: scripts/train_from_data.sh [-o OUTPUT_FILE]
|
|
|
|
Trains the model on complete session triplets found under data/{motion,heart_rate,labels}/.
|
|
EOF
|
|
exit 0
|
|
fi
|
|
|
|
if [[ ${1:-} == "-o" || ${1:-} == "--output" ]]; then
|
|
OUTPUT_FILE="${2:?Missing output file after -o/--output}"
|
|
MODEL_FILE="${OUTPUT_FILE}"
|
|
fi
|
|
|
|
shopt -s nullglob
|
|
|
|
checkpoint_class_count() {
|
|
python - "$1" <<'PY'
|
|
from pathlib import Path
|
|
import sys
|
|
import torch
|
|
|
|
payload = torch.load(Path(sys.argv[1]), weights_only=True)
|
|
state_dict = payload.get("model_state_dict", payload)
|
|
print(state_dict["net.2.bias"].shape[0])
|
|
PY
|
|
}
|
|
|
|
motions=("${DATA_DIR}/motion"/*_acceleration.txt)
|
|
|
|
subject_ids=()
|
|
|
|
for motion_file in "${motions[@]}"; do
|
|
sample_id="$(basename "${motion_file}" _acceleration.txt)"
|
|
heartrate_file="${DATA_DIR}/heart_rate/${sample_id}_heartrate.txt"
|
|
label_file="${DATA_DIR}/labels/${sample_id}_labeled_sleep.txt"
|
|
|
|
if [[ -f "${heartrate_file}" && -f "${label_file}" ]]; then
|
|
subject_ids+=("${sample_id}")
|
|
fi
|
|
done
|
|
|
|
if [[ ${#subject_ids[@]} -eq 0 ]]; then
|
|
printf 'No complete data triplets found under %s\n' "${DATA_DIR}" >&2
|
|
exit 1
|
|
fi
|
|
|
|
mapfile -t subject_ids < <(printf '%s\n' "${subject_ids[@]}" | sort -u)
|
|
|
|
printf 'Training on %d subject(s)\n' "${#subject_ids[@]}"
|
|
|
|
export PYTHONPATH="${ROOT_DIR}/src${PYTHONPATH:+:${PYTHONPATH}}"
|
|
|
|
for subject_id in "${subject_ids[@]}"; do
|
|
acceleration_file="${DATA_DIR}/motion/${subject_id}_acceleration.txt"
|
|
heartrate_file="${DATA_DIR}/heart_rate/${subject_id}_heartrate.txt"
|
|
label_file="${DATA_DIR}/labels/${subject_id}_labeled_sleep.txt"
|
|
resume_args=()
|
|
|
|
printf 'Training subject %s\n' "${subject_id}"
|
|
|
|
if [[ -f "${MODEL_FILE}" ]]; then
|
|
checkpoint_classes="$(checkpoint_class_count "${MODEL_FILE}")"
|
|
if [[ "${checkpoint_classes}" == "3" ]]; then
|
|
resume_args=(-i "${MODEL_FILE}")
|
|
else
|
|
printf 'Skipping checkpoint for subject %s: model classes=%s, expected=3\n' "${subject_id}" "${checkpoint_classes}"
|
|
fi
|
|
fi
|
|
|
|
args=(python -m sleep_detection train -a "${acceleration_file}" -H "${heartrate_file}" -l "${label_file}" -o "${MODEL_FILE}")
|
|
if [[ ${#resume_args[@]} -gt 0 ]]; then
|
|
args+=("${resume_args[@]}")
|
|
fi
|
|
|
|
"${args[@]}"
|
|
done
|
|
|
|
printf 'Saved trained model to %s\n' "${MODEL_FILE}"
|