File size: 1,959 Bytes
dd253d4 | 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 | #!/bin/bash
# Generic 3-PoinTr flow-backbone trainer. Set via env: DATA_ROOT, TASK, MODEL_NAME,
# SUBSAMPLE (<= min_N), USE_LABELS (0/1 → --use_point_labels for 2-object), GPU, EPOCHS.
set -e
DATA_ROOT="${DATA_ROOT:?set DATA_ROOT}"
TASK="${TASK:?set TASK}"
MODEL_NAME="${MODEL_NAME:?set MODEL_NAME}"
SUBSAMPLE="${SUBSAMPLE:-150}"
USE_LABELS="${USE_LABELS:-0}"
SPLIT_NAME="${TASK}_flow"
EPOCHS="${EPOCHS:-5000}"
BATCH_SIZE="${BATCH_SIZE:-4}"
GPU="${GPU:-0}"
cd /home/hongyi/PWM/reward_model/3PoinTr_repo
source /home/hongyi/miniconda3/etc/profile.d/conda.sh
conda activate pointr
TOTAL=$(ls -d "${DATA_ROOT}/no_actions/${TASK}"/[0-9][0-9][0-9][0-9] 2>/dev/null | wc -l)
[ "$TOTAL" -lt 5 ] && { echo "ERROR: only ${TOTAL} episodes built"; exit 1; }
N_VAL=$((TOTAL / 10)); [ "$N_VAL" -lt 1 ] && N_VAL=1
N_TRAIN=$((TOTAL - N_VAL))
echo "[launch] ${TASK}: ${TOTAL} eps -> train=${N_TRAIN} val=${N_VAL} subsample=${SUBSAMPLE} labels=${USE_LABELS} -> ${MODEL_NAME}"
TASK_SPLITS="{\"${TASK}\": [0, 0, 0, ${N_TRAIN}, ${N_VAL}, 0]}"
python scripts_training/splitter.py \
--dataset "${DATA_ROOT}/" --task-splits "${TASK_SPLITS}" \
--out-dir "${DATA_ROOT}/" --name "${SPLIT_NAME}"
SPLIT_FILE="${DATA_ROOT}/${SPLIT_NAME}.json"
LABEL_FLAG=""; [ "$USE_LABELS" = "1" ] && LABEL_FLAG="--use_point_labels"
mkdir -p /home/hongyi/PWM/reward_model/3pointr_models
CUDA_VISIBLE_DEVICES=$GPU python scripts_training/pretrain_flow.py \
--device cuda:0 --epochs "${EPOCHS}" --batch_size "${BATCH_SIZE}" --lr 1e-4 \
--arch_dim 256 --arch_depth 2 --arch_heads 4 --arch_query_depth 1 --arch_query_heads 1 \
--subsample "${SUBSAMPLE}" --point_dim 2 ${LABEL_FLAG} --no_img_cond \
--dataset "${DATA_ROOT}/" --model_name "${MODEL_NAME}" --tasks "${TASK}" \
--save_dir "/home/hongyi/PWM/reward_model/3pointr_models" --split_file "${SPLIT_FILE}" \
--flow_data_only --eval_every 50 --save_vis_every 250 --weight_decay 1e-4 --seed 42 \
--num_workers 4 --flow_loss_w 10.0
|