File size: 1,394 Bytes
275b5a1 923fdff 275b5a1 923fdff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | #!/usr/bin/env bash
set -euo pipefail
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd -- "${SCRIPT_DIR}/.." && pwd)"
PYTHON_BIN="${PYTHON_BIN:-python}"
CONFIG="${CONFIG:-${SCRIPT_DIR}/configs/stead_empty8_local.yaml}"
TRAIN_H5="${PROJECT_ROOT}/Train/data/processed/from_raw/stead/train/stead_100hz_60s_train.h5"
NUM_GPUS="${NUM_GPUS:-8}"
while [[ $# -gt 0 ]]; do
case "$1" in
--num-gpus)
[[ $# -ge 2 ]] || { echo "--num-gpus requires a value" >&2; exit 2; }
NUM_GPUS="$2"
shift 2
;;
--num-gpus=*)
NUM_GPUS="${1#*=}"
shift
;;
*)
echo "Unknown argument: $1" >&2
echo "Usage: bash Train/run_train_local.sh [--num-gpus N]" >&2
exit 2
;;
esac
done
[[ "${NUM_GPUS}" =~ ^[1-9][0-9]*$ ]] || { echo "NUM_GPUS must be a positive integer" >&2; exit 2; }
cd "${PROJECT_ROOT}"
[[ -f "${TRAIN_H5}" ]] || {
echo "Missing processed STEAD data: ${TRAIN_H5}" >&2
echo "Run: bash Train/run_download_prepare.sh" >&2
exit 1
}
export PYTHONUNBUFFERED=1
echo "Training with ${NUM_GPUS} GPU(s); global batch size is read from ${CONFIG}"
if [[ "${NUM_GPUS}" -eq 1 ]]; then
exec "${PYTHON_BIN}" Train/training/train.py --config "${CONFIG}"
fi
exec "${PYTHON_BIN}" -m torch.distributed.run \
--standalone \
--nproc_per_node="${NUM_GPUS}" \
Train/training/train.py --config "${CONFIG}"
|