pose6daug-scripts / train /smoke_test.sh

pose6daug augmentation / dataset / training / eval scripts

5d08972 verified about 1 month ago

3.14 kB

#!/usr/bin/env bash
# Verifies that GR00T 1.5 fine-tuning runs from THIS bundle: imports, dataset load,
# then a REAL 20-step single-GPU fine-tune on the 150-episode augmented dataset
# (per-GPU batch 2, no LoRA, same code path as the full run) writing a checkpoint
# into runs/smoke_<timestamp>/.  Needs one GPU with ~25 GB free.
#   ./smoke_test.sh <gpu_index>          e.g. ./smoke_test.sh 3
set -euo pipefail
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"; B="$(dirname "$HERE")"
GPU="${1:?usage: smoke_test.sh <gpu_index>}"
PY=/data/nvidia/gripper_augmentator/conda-envs/mygr00t/bin/python
DS="${DATASET_PATH:-/lp-dev/jonghoon/robocasa_full/gr00t_views/PickPlaceCounterToCabinet_150eps_augmented_groot15_v3_100_clean}"
CK="${BASE_CKPT:-/lp-dev/jonghoon/myGR00T_outputs/pnpcountertocab_all502_gbs64_wandb_60k_save5k_20260427_193901/checkpoint-60000}"
pass(){ echo "  PASS  $*"; }; fail(){ echo "  FAIL  $*"; exit 1; }
export PYTHONPATH="$B/myGR00T" HF_HOME=/lp-dev/jonghoon/.cache/huggingface HF_HUB_OFFLINE=1 NO_ALBUMENTATIONS_UPDATE=1

echo "== static =="
[ -x "$PY" ] && pass "env python $PY" || fail "mygr00t env missing"
(cd "$B/myGR00T" && "$PY" -c "import gr00t, os; assert gr00t.__file__.startswith('$B/'), gr00t.__file__; from gr00t.experiment.data_config import DATA_CONFIG_MAP; assert 'single_panda_gripper' in DATA_CONFIG_MAP") && pass "gr00t imports from the bundle, data config single_panda_gripper registered" || fail "gr00t import"
[ -f "$DS/meta/modality.json" ] && pass "dataset $DS" || fail "dataset missing"
[ -f "$CK/model.safetensors.index.json" ] && pass "base checkpoint $CK" || fail "checkpoint missing"
(cd "$B/myGR00T" && "$PY" - <<PYEOF
from gr00t.data.dataset import LeRobotSingleDataset
from gr00t.data.schema import EmbodimentTag
from gr00t.experiment.data_config import DATA_CONFIG_MAP
dc = DATA_CONFIG_MAP["single_panda_gripper"]
ds = LeRobotSingleDataset("$DS", dc.modality_config(), transforms=dc.transform(), embodiment_tag=EmbodimentTag("new_embodiment"), video_backend="decord")
s = ds[0]; print({k: tuple(v.shape) for k, v in s.items() if hasattr(v, "shape")}); assert len(ds) > 1000
PYEOF
) && pass "LeRobotSingleDataset loads a sample (videos decoded)" || fail "dataset load"

echo "== 20-step fine-tune on GPU $GPU =="
RUN=smoke_$(date -u +%Y%m%d_%H%M%S); OUT="$B/runs/$RUN"
CUDA_VISIBLE_DEVICES="$GPU" GPUS="$GPU" NUM_GPUS=1 PER_GPU_BATCH=2 MAX_STEPS=20 SAVE_STEPS=20 DATALOADER_NUM_WORKERS=4 \
  REPORT_TO=tensorboard RUN_NAME="$RUN" OUTPUT_DIR="$OUT" DATASET_PATH="$DS" BASE_CKPT="$CK" \
  bash "$HERE/train_groot15_single_dataset.sh" > "$B/logs/$RUN.log" 2>&1 || { tail -30 "$B/logs/$RUN.log"; fail "fine-tune exited non-zero (log: $B/logs/$RUN.log)"; }
[ -f "$OUT/checkpoint-20/model.safetensors.index.json" ] || [ -f "$OUT/checkpoint-20/model.safetensors" ] && pass "checkpoint-20 written to $OUT" || fail "no checkpoint-20 in $OUT"
grep -q "train_loss\|'loss'" "$B/logs/$RUN.log" && pass "loss logged: $(grep -o "'loss': [0-9.]*" "$B/logs/$RUN.log" | tail -1)" || pass "run finished (no loss line found in log, check $B/logs/$RUN.log)"
echo "ALL PASSED  (log: $B/logs/$RUN.log)"

	#!/usr/bin/env bash
	# Verifies that GR00T 1.5 fine-tuning runs from THIS bundle: imports, dataset load,
	# then a REAL 20-step single-GPU fine-tune on the 150-episode augmented dataset
	# (per-GPU batch 2, no LoRA, same code path as the full run) writing a checkpoint
	# into runs/smoke_<timestamp>/. Needs one GPU with ~25 GB free.
	# ./smoke_test.sh <gpu_index> e.g. ./smoke_test.sh 3
	set -euo pipefail
	HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"; B="$(dirname "$HERE")"
	GPU="${1:?usage: smoke_test.sh <gpu_index>}"
	PY=/data/nvidia/gripper_augmentator/conda-envs/mygr00t/bin/python
	DS="${DATASET_PATH:-/lp-dev/jonghoon/robocasa_full/gr00t_views/PickPlaceCounterToCabinet_150eps_augmented_groot15_v3_100_clean}"
	CK="${BASE_CKPT:-/lp-dev/jonghoon/myGR00T_outputs/pnpcountertocab_all502_gbs64_wandb_60k_save5k_20260427_193901/checkpoint-60000}"
	pass(){ echo " PASS $"; }; fail(){ echo " FAIL $"; exit 1; }
	export PYTHONPATH="$B/myGR00T" HF_HOME=/lp-dev/jonghoon/.cache/huggingface HF_HUB_OFFLINE=1 NO_ALBUMENTATIONS_UPDATE=1

	echo "== static =="
	[ -x "$PY" ] && pass "env python $PY" \|\| fail "mygr00t env missing"
	(cd "$B/myGR00T" && "$PY" -c "import gr00t, os; assert gr00t.__file__.startswith('$B/'), gr00t.__file__; from gr00t.experiment.data_config import DATA_CONFIG_MAP; assert 'single_panda_gripper' in DATA_CONFIG_MAP") && pass "gr00t imports from the bundle, data config single_panda_gripper registered" \|\| fail "gr00t import"
	[ -f "$DS/meta/modality.json" ] && pass "dataset $DS" \|\| fail "dataset missing"
	[ -f "$CK/model.safetensors.index.json" ] && pass "base checkpoint $CK" \|\| fail "checkpoint missing"
	(cd "$B/myGR00T" && "$PY" - <<PYEOF
	from gr00t.data.dataset import LeRobotSingleDataset
	from gr00t.data.schema import EmbodimentTag
	from gr00t.experiment.data_config import DATA_CONFIG_MAP
	dc = DATA_CONFIG_MAP["single_panda_gripper"]
	ds = LeRobotSingleDataset("$DS", dc.modality_config(), transforms=dc.transform(), embodiment_tag=EmbodimentTag("new_embodiment"), video_backend="decord")
	s = ds[0]; print({k: tuple(v.shape) for k, v in s.items() if hasattr(v, "shape")}); assert len(ds) > 1000
	PYEOF
	) && pass "LeRobotSingleDataset loads a sample (videos decoded)" \|\| fail "dataset load"

	echo "== 20-step fine-tune on GPU $GPU =="
	RUN=smoke_$(date -u +%Y%m%d_%H%M%S); OUT="$B/runs/$RUN"
	CUDA_VISIBLE_DEVICES="$GPU" GPUS="$GPU" NUM_GPUS=1 PER_GPU_BATCH=2 MAX_STEPS=20 SAVE_STEPS=20 DATALOADER_NUM_WORKERS=4 \
	REPORT_TO=tensorboard RUN_NAME="$RUN" OUTPUT_DIR="$OUT" DATASET_PATH="$DS" BASE_CKPT="$CK" \
	bash "$HERE/train_groot15_single_dataset.sh" > "$B/logs/$RUN.log" 2>&1 \|\| { tail -30 "$B/logs/$RUN.log"; fail "fine-tune exited non-zero (log: $B/logs/$RUN.log)"; }
	[ -f "$OUT/checkpoint-20/model.safetensors.index.json" ] \|\| [ -f "$OUT/checkpoint-20/model.safetensors" ] && pass "checkpoint-20 written to $OUT" \|\| fail "no checkpoint-20 in $OUT"
	grep -q "train_loss\\|'loss'" "$B/logs/$RUN.log" && pass "loss logged: $(grep -o "'loss': [0-9.]*" "$B/logs/$RUN.log" \| tail -1)" \|\| pass "run finished (no loss line found in log, check $B/logs/$RUN.log)"
	echo "ALL PASSED (log: $B/logs/$RUN.log)"