Initial RRoG GNN runner

8 files changed, 236 insertions, 0 deletions

diff --git a/scripts/collect_results.sh b/scripts/collect_results.sh
new file mode 100755
index 0000000..360f05b
--- /dev/null
+++ b/scripts/collect_results.sh
@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "${ROOT_DIR}"
+export PYTHONPATH="${ROOT_DIR}:${PYTHONPATH:-}"
+
+mkdir -p summaries
+python3 -m rrog.cli zinc-results --epochs "${ZINC_EPOCHS:-200}" | tee summaries/zinc_cycle56.md
+python3 -m rrog.cli results --epochs "${OGB_EPOCHS:-100}" | tee summaries/ogb_graphprop.md
diff --git a/scripts/run_ogb_mol_all_tasks.sh b/scripts/run_ogb_mol_all_tasks.sh
new file mode 100755
index 0000000..b191d79
--- /dev/null
+++ b/scripts/run_ogb_mol_all_tasks.sh
@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "${ROOT_DIR}"
+
+DEVICE="${DEVICE:-cuda:1}"
+EPOCHS="${EPOCHS:-100}"
+SEED="${SEED:-0}"
+TASKS="${TASKS:-ogbg-molhiv ogbg-molbbbp ogbg-molbace ogbg-moltox21 ogbg-molclintox ogbg-molsider ogbg-molesol ogbg-molfreesolv ogbg-mollipo}"
+
+mkdir -p logs
+for task in ${TASKS}; do
+  echo "[task] ${task}"
+  TASK="${task}" DEVICE="${DEVICE}" EPOCHS="${EPOCHS}" SEED="${SEED}" \
+    ./scripts/run_ogb_mol_task_full.sh 2>&1 | tee "logs/${task}_${SEED}.log"
+done
diff --git a/scripts/run_ogb_mol_task_full.sh b/scripts/run_ogb_mol_task_full.sh
new file mode 100755
index 0000000..b25bff3
--- /dev/null
+++ b/scripts/run_ogb_mol_task_full.sh
@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "${ROOT_DIR}"
+export PYTHONPATH="${ROOT_DIR}:${PYTHONPATH:-}"
+
+TASK="${TASK:-ogbg-molhiv}"
+DEVICE="${DEVICE:-cuda:1}"
+EPOCHS="${EPOCHS:-100}"
+SEED="${SEED:-0}"
+HIDDEN="${HIDDEN:-128}"
+VIEWS="${VIEWS:-gin gine gcn graphsage gatv2 graphconv transformer pna gen film resgated tag sgc cheb arma mf appnp}"
+
+mkdir -p runs logs
+
+result_path() {
+  local view="$1"
+  local compute="$2"
+  local t="$3"
+  local ns="$4"
+  echo "runs/${TASK}_${view}_${compute}_T${t}_ns${ns}_h${HIDDEN}_e${EPOCHS}_s${SEED}.json"
+}
+
+run_cell() {
+  local view="$1"
+  local compute="$2"
+  local t="$3"
+  local ns="$4"
+  local out
+  out="$(result_path "${view}" "${compute}" "${t}" "${ns}")"
+  if [[ -f "${out}" ]]; then
+    echo "[skip] ${out}"
+    return
+  fi
+  echo "[run] ${TASK} view=${view} compute=${compute} T=${t} ns=${ns} device=${DEVICE}"
+  python3 -m rrog.cli run \
+    --task "${TASK}" \
+    --view "${view}" \
+    --compute "${compute}" \
+    --epochs "${EPOCHS}" \
+    --hidden "${HIDDEN}" \
+    --T "${t}" \
+    --n_sup "${ns}" \
+    --seed "${SEED}" \
+    --device "${DEVICE}"
+}
+
+for view in ${VIEWS}; do
+  run_cell "${view}" classic 0 1
+  run_cell "${view}" fixed-rrog 3 3
+done
+
+python3 -m rrog.cli results --epochs "${EPOCHS}"
diff --git a/scripts/run_smoke.sh b/scripts/run_smoke.sh
new file mode 100755
index 0000000..6365cec
--- /dev/null
+++ b/scripts/run_smoke.sh
@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "${ROOT_DIR}"
+export PYTHONPATH="${ROOT_DIR}:${PYTHONPATH:-}"
+
+DEVICE="${DEVICE:-cuda:0}"
+mkdir -p runs logs
+
+python3 -m rrog.cli run \
+  --task ogbg-molhiv --view gin --compute classic \
+  --epochs 1 --hidden 32 --bs 64 --seed 991 --device "${DEVICE}" \
+  --max_train_batches 2 --max_eval_batches 2
+
+python3 -m rrog.cli run \
+  --task ogbg-molhiv --view gin --compute fixed-rrog \
+  --epochs 1 --hidden 32 --bs 64 --T 1 --n_sup 2 --seed 992 --device "${DEVICE}" \
+  --max_train_batches 2 --max_eval_batches 2
diff --git a/scripts/run_two_a6000.sh b/scripts/run_two_a6000.sh
new file mode 100755
index 0000000..8d9851f
--- /dev/null
+++ b/scripts/run_two_a6000.sh
@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "${ROOT_DIR}"
+export PYTHONPATH="${ROOT_DIR}:${PYTHONPATH:-}"
+
+ZINC_DEVICE="${ZINC_DEVICE:-cuda:0}"
+OGB_DEVICE="${OGB_DEVICE:-cuda:1}"
+OGB_TASK="${OGB_TASK:-ogbg-molhiv}"
+ZINC_EPOCHS="${ZINC_EPOCHS:-200}"
+OGB_EPOCHS="${OGB_EPOCHS:-100}"
+SEED="${SEED:-0}"
+
+mkdir -p runs logs
+
+echo "[launch] ZINC-cycle56 on ${ZINC_DEVICE}"
+DEVICE="${ZINC_DEVICE}" EPOCHS="${ZINC_EPOCHS}" SEED="${SEED}" \
+  ./scripts/run_zinc_cycle56_full.sh > "logs/zinc_cycle56_${SEED}.log" 2>&1 &
+zinc_pid=$!
+
+echo "[launch] ${OGB_TASK} on ${OGB_DEVICE}"
+TASK="${OGB_TASK}" DEVICE="${OGB_DEVICE}" EPOCHS="${OGB_EPOCHS}" SEED="${SEED}" \
+  ./scripts/run_ogb_mol_task_full.sh > "logs/${OGB_TASK}_${SEED}.log" 2>&1 &
+ogb_pid=$!
+
+echo "[pids] zinc=${zinc_pid} ogb=${ogb_pid}"
+wait "${zinc_pid}"
+wait "${ogb_pid}"
+
+echo "[done] collecting summaries"
+./scripts/collect_results.sh
diff --git a/scripts/run_zinc_cycle56_full.sh b/scripts/run_zinc_cycle56_full.sh
new file mode 100755
index 0000000..151a51e
--- /dev/null
+++ b/scripts/run_zinc_cycle56_full.sh
@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "${ROOT_DIR}"
+export PYTHONPATH="${ROOT_DIR}:${PYTHONPATH:-}"
+
+DEVICE="${DEVICE:-cuda:0}"
+EPOCHS="${EPOCHS:-200}"
+SEED="${SEED:-0}"
+VIEWS="${VIEWS:-gin gine gcn graphsage gatv2 graphconv transformer pna gen film resgated tag sgc cheb arma mf appnp}"
+
+mkdir -p runs logs
+
+result_path() {
+  local view="$1"
+  local t="$2"
+  local ns="$3"
+  local view_tag=""
+  if [[ "${view}" != "gin" ]]; then
+    view_tag="_${view}"
+  fi
+  echo "runs/rec_rrog${view_tag}_full_sig0.0_K1_none_T${t}_ns${ns}_trace_s${SEED}.json"
+}
+
+run_cell() {
+  local view="$1"
+  local compute="$2"
+  local t="$3"
+  local ns="$4"
+  local out
+  out="$(result_path "${view}" "${t}" "${ns}")"
+  if [[ -f "${out}" ]]; then
+    echo "[skip] ${out}"
+    return
+  fi
+  echo "[run] zinc-cycle56 view=${view} compute=${compute} T=${t} ns=${ns} device=${DEVICE}"
+  python3 -m rrog.cli run \
+    --task zinc-cycle56 \
+    --view "${view}" \
+    --compute "${compute}" \
+    --epochs "${EPOCHS}" \
+    --T "${t}" \
+    --n_sup "${ns}" \
+    --seed "${SEED}" \
+    --device "${DEVICE}"
+}
+
+for view in ${VIEWS}; do
+  run_cell "${view}" classic 0 1
+  run_cell "${view}" fixed-rrog 1 3
+done
+
+python3 -m rrog.cli zinc-results --epochs "${EPOCHS}"
diff --git a/scripts/setup_and_run_two_a6000.sh b/scripts/setup_and_run_two_a6000.sh
new file mode 100755
index 0000000..ec4e3da
--- /dev/null
+++ b/scripts/setup_and_run_two_a6000.sh
@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "${ROOT_DIR}"
+
+if [[ "${SKIP_SETUP:-0}" != "1" ]]; then
+  ./scripts/setup_env.sh
+fi
+
+if [[ -d "${VENV_DIR:-.venv}" ]]; then
+  source "${VENV_DIR:-.venv}/bin/activate"
+fi
+
+./scripts/run_two_a6000.sh
diff --git a/scripts/setup_env.sh b/scripts/setup_env.sh
new file mode 100755
index 0000000..66a94c8
--- /dev/null
+++ b/scripts/setup_env.sh
@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "${ROOT_DIR}"
+
+PYTHON_BIN="${PYTHON_BIN:-python3}"
+VENV_DIR="${VENV_DIR:-.venv}"
+TORCH_INDEX_URL="${TORCH_INDEX_URL:-https://download.pytorch.org/whl/cu124}"
+
+if [[ ! -d "${VENV_DIR}" ]]; then
+  "${PYTHON_BIN}" -m venv "${VENV_DIR}"
+fi
+
+source "${VENV_DIR}/bin/activate"
+python -m pip install --upgrade pip wheel setuptools
+
+if ! python - <<'PY' >/dev/null 2>&1
+import torch
+assert torch.cuda.is_available() or True
+PY
+then
+  python -m pip install torch --index-url "${TORCH_INDEX_URL}"
+fi
+
+python -m pip install -r requirements.txt
+
+python - <<'PY'
+import torch
+import torch_geometric
+import ogb
+print("torch", torch.__version__, "cuda_available", torch.cuda.is_available())
+print("torch_geometric", torch_geometric.__version__)
+print("ogb", ogb.__version__)
+PY