File size: 4,074 Bytes
a26fd8c | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | #!/bin/bash
# Noisy-latent Sudoku trainer. Works from either layout:
# Sudoku_superposition/ (HF: code/ + data/)
# sudoku-wavecurriculum/ (GitHub branch; set SUDOKU_DATA_DIR)
#
# Usage:
# bash run_experiment.sh
# SUDOKU_INSTANCE_DIR=./datasets_superposition bash run_experiment.sh
# SUDOKU_DATA_DIR=/path/to/data bash run_experiment.sh
set -euo pipefail
die() { echo "ERROR: $*" >&2; exit 1; }
need_file() { [ -s "$1" ] || die "missing or empty: $1"; }
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
if [ -d "${HERE}/train" ] && [ -f "${HERE}/train/main.py" ]; then
CODE="${HERE}"
ROOT="$(cd "${HERE}/.." && pwd)"
elif [ -d "${HERE}/code/train" ] && [ -f "${HERE}/code/train/main.py" ]; then
CODE="${HERE}/code"
ROOT="${HERE}"
elif [ -d "${HERE}/../data" ] && [ -f "${HERE}/train/main.py" ]; then
CODE="${HERE}"
ROOT="$(cd "${HERE}/.." && pwd)"
else
die "cannot find train/main.py next to this script or in ./code/"
fi
# HF download puts arrays in data/. New instance rebuilds go in datasets_superposition/.
DATA="${SUDOKU_DATA_DIR:-}"
if [ -z "${DATA}" ]; then
if [ -d "${ROOT}/data" ]; then
DATA="${ROOT}/data"
elif [ -d "${HERE}/data" ]; then
DATA="${HERE}/data"
else
die "set SUDOKU_DATA_DIR to the folder with the 11 npy/npz files"
fi
fi
DATA="$(cd "${DATA}" && pwd)"
INST="${SUDOKU_INSTANCE_DIR:-}"
if [ -z "${INST}" ]; then
if [ -s "${ROOT}/datasets_superposition/train_assignments.npy" ] \
&& [ -s "${ROOT}/datasets_superposition/train_starts.npy" ] \
&& [ -s "${ROOT}/datasets_superposition/train_counts.npy" ]; then
INST="${ROOT}/datasets_superposition"
echo "[run] using rebuilt instances in ${INST}"
else
INST="${DATA}"
fi
fi
INST="$(cd "${INST}" && pwd)"
export SUDOKU_TRAIN_PATH="${SUDOKU_TRAIN_PATH:-${DATA}/train_sudoku_puzzles.npy}"
export SUDOKU_TEST_PATH="${SUDOKU_TEST_PATH:-${DATA}/test_sudoku_puzzles.npy}"
export SUDOKU_TRAIN_CAND="${SUDOKU_TRAIN_CAND:-${DATA}/train_cand_masks.npy}"
export SUDOKU_TEST_CAND="${SUDOKU_TEST_CAND:-${DATA}/test_cand_masks.npy}"
export SUDOKU_INSTANCE_DIR="${INST}"
export SUDOKU_GATE_Q="${SUDOKU_GATE_Q:-${DATA}/test_bigq300.npz}"
need_file "${SUDOKU_TRAIN_PATH}"
need_file "${SUDOKU_TEST_PATH}"
need_file "${SUDOKU_TRAIN_CAND}"
need_file "${SUDOKU_TEST_CAND}"
need_file "${SUDOKU_GATE_Q}"
for split in train test; do
need_file "${INST}/${split}_assignments.npy"
need_file "${INST}/${split}_starts.npy"
need_file "${INST}/${split}_counts.npy"
done
export SUDOKU_LATENT_SLOTS="${SUDOKU_LATENT_SLOTS:-12}"
export SUDOKU_AUX_WEIGHT="${SUDOKU_AUX_WEIGHT:-0}"
export SUDOKU_LEAK_WEIGHT="${SUDOKU_LEAK_WEIGHT:-0}"
export SUDOKU_LATENT_NOISE="${SUDOKU_LATENT_NOISE:-1}"
export SUDOKU_MC_SAMPLES="${SUDOKU_MC_SAMPLES:-8}"
export SUDOKU_MC_OBJ="${SUDOKU_MC_OBJ:-mean}"
export SUDOKU_NOISE_LOGSIGMA0="${SUDOKU_NOISE_LOGSIGMA0:--2.5}"
export SUDOKU_START_STAGE="${SUDOKU_START_STAGE:-1}"
export SUDOKU_MAX_STAGE="${SUDOKU_MAX_STAGE:-12}"
export SUDOKU_MINIBATCH="${SUDOKU_MINIBATCH:-64}"
export SUDOKU_EVAL_EVERY="${SUDOKU_EVAL_EVERY:-2000}"
export SUDOKU_MAX_STEPS="${SUDOKU_MAX_STEPS:-2500000}"
PY="${PYTHON:-}"
if [ -z "${PY}" ]; then
if command -v python >/dev/null 2>&1; then
PY=python
elif command -v python3 >/dev/null 2>&1; then
PY=python3
else
die "python not on PATH"
fi
fi
# Trainer imports these at module load even when wandb is off.
"${PY}" -c "import jax, flax, optax, ml_collections, tensorflow, wandb, absl" \
|| die "install: jax flax optax ml_collections clu absl-py tensorflow wandb"
WORKDIR="${WORKDIR:-${ROOT}/runs/noise}"
EXP_NAME="${SUDOKU_EXP_NAME:-noise}"
mkdir -p "${WORKDIR}"
echo "[run] code=${CODE}"
echo "[run] puzzles/masks/q=${DATA}"
echo "[run] instances=${INST}"
echo "[run] workdir=${WORKDIR}"
echo "[run] noise=${SUDOKU_LATENT_NOISE} M=${SUDOKU_MC_SAMPLES} K=${SUDOKU_LATENT_SLOTS} start=${SUDOKU_START_STAGE}"
cd "${CODE}"
[ -f train/main.py ] || die "train/main.py not in ${CODE}"
exec "${PY}" -m train.main --workdir="${WORKDIR}" --exp_name="${EXP_NAME}"
|