Add causal person track GRU residual experiment

2026-09-07 00:30:23 +08:00
parent 9fc4577aea
commit 4aed4697a5
15 changed files with 1185 additions and 43 deletions
@@ -0,0 +1,177 @@
+#!/usr/bin/env bash
+# Foreground matched EasyCom pooled AV -> causal anonymous-track GRU residual.
+# Usage: bash run_easycom_track_gru_residual.sh [all|pooled|track|test] [seed ...]
+set -euo pipefail
+cd "$(dirname "$0")"
+
+FEATURE_ROOT=${FEATURE_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack}
+TARGET_ROOT=${TARGET_ROOT:-/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame}
+LOCAL_ROOT=${LOCAL_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/local_qwen_1s_encoded_tf4576}
+CHECKPOINT_ROOT=${CHECKPOINT_ROOT:-checkpoints/easycom_track_gru_residual}
+RESULT_ROOT=${RESULT_ROOT:-results/easycom_track_gru_residual}
+TEST_RESULT_ROOT=${TEST_RESULT_ROOT:-results/easycom_track_gru_residual_test}
+EPOCHS=${EPOCHS:-4}
+WARMUP_EPOCHS=${WARMUP_EPOCHS:-1.6}
+LOCAL_PROJECTION_DIM=${LOCAL_PROJECTION_DIM:-256}
+PERSON_ATTENTION_DIM=${PERSON_ATTENTION_DIM:-256}
+PERSON_ATTENTION_HEADS=${PERSON_ATTENTION_HEADS:-4}
+PERSON_HISTORY_FRAMES=${PERSON_HISTORY_FRAMES:-16}
+PERSON_TRACK_MAX_GAP_SECONDS=${PERSON_TRACK_MAX_GAP_SECONDS:-1.0}
+SUMMARIZE=${SUMMARIZE:-1}
+read -r -a SESSION_ARGS <<< "${SESSIONS:-4 5 6 7 8 9 10 11 12}"
+read -r -a DEV_SESSION_ARGS <<< "${DEV_SESSIONS:-12}"
+read -r -a SCREEN_SESSION_ARGS <<< "${SCREEN_SESSIONS:-12}"
+export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4}
+export MKL_NUM_THREADS=${MKL_NUM_THREADS:-4}
+export OPENBLAS_NUM_THREADS=${OPENBLAS_NUM_THREADS:-4}
+
+variant=${1:-all}
+if (( $# )); then shift; fi
+if (( $# )); then seeds=("$@"); else seeds=(0 1 2); fi
+case "$variant" in
+  all) variants=(pooled track) ;;
+  pooled|track) variants=("$variant") ;;
+  test) variants=(pooled track) ;;
+  *) echo "Expected all, pooled, track, or test" >&2; exit 2 ;;
+esac
+
+mkdir -p "$CHECKPOINT_ROOT" "$RESULT_ROOT"
+
+infer_checkpoint() {
+  local checkpoint=$1 result=$2
+  shift 2
+  if [[ -e "$result/metrics.json" ]]; then
+    echo "Refusing to overwrite $result/metrics.json" >&2
+    return 1
+  fi
+  python infer_turn_head.py \
+    --feature-root "$FEATURE_ROOT" \
+    --target-root "$TARGET_ROOT" \
+    --local-feature-root "$LOCAL_ROOT" \
+    --checkpoint "$checkpoint" \
+    --output-dir "$result" \
+    --sessions "$@" \
+    --expected-mode open-stream-replay \
+    --device cuda:0
+}
+
+train_and_infer_pooled() {
+  local seed=$1 checkpoint result
+  checkpoint="$CHECKPOINT_ROOT/pooled_seed${seed}.pth"
+  result="$RESULT_ROOT/pooled_seed${seed}"
+  if [[ -e "$checkpoint" || -e "$result/metrics.json" ]]; then
+    echo "Refusing to overwrite $checkpoint or $result" >&2
+    return 1
+  fi
+  python train_turn_head.py \
+    --feature-root "$FEATURE_ROOT" \
+    --target-root "$TARGET_ROOT" \
+    --local-feature-root "$LOCAL_ROOT" \
+    --local-modalities av \
+    --local-fusion pooled \
+    --local-projection-dim "$LOCAL_PROJECTION_DIM" \
+    --output "$checkpoint" \
+    --sessions "${SESSION_ARGS[@]}" \
+    --dev-sessions "${DEV_SESSION_ARGS[@]}" \
+    --epochs "$EPOCHS" \
+    --batch-size 512 \
+    --learning-rate 7e-5 \
+    --weight-decay 5e-5 \
+    --warmup-epochs "$WARMUP_EPOCHS" \
+    --warmup-factor 0.3 \
+    --dropout 0.2 \
+    --loss lstr-mce \
+    --transition-objective none \
+    --layer-fusion concat-mlp \
+    --expected-mode open-stream-replay \
+    --seed "$seed" \
+    --device cuda:0
+  # Screen architecture/hyperparameters only on the held-out training session.
+  infer_checkpoint "$checkpoint" "$result" "${SCREEN_SESSION_ARGS[@]}"
+}
+
+train_and_infer_track() {
+  local seed=$1 base checkpoint result
+  base="$CHECKPOINT_ROOT/pooled_seed${seed}.pth"
+  checkpoint="$CHECKPOINT_ROOT/track_gru_seed${seed}.pth"
+  result="$RESULT_ROOT/track_gru_seed${seed}"
+  if [[ ! -s "$base" ]]; then
+    echo "Missing matched pooled checkpoint: $base; run variant pooled or all first" >&2
+    return 1
+  fi
+  if [[ -e "$checkpoint" || -e "$result/metrics.json" ]]; then
+    echo "Refusing to overwrite $checkpoint or $result" >&2
+    return 1
+  fi
+  python train_turn_head.py \
+    --feature-root "$FEATURE_ROOT" \
+    --target-root "$TARGET_ROOT" \
+    --local-feature-root "$LOCAL_ROOT" \
+    --local-modalities av \
+    --local-fusion track-gru-residual \
+    --local-projection-dim "$LOCAL_PROJECTION_DIM" \
+    --person-attention-dim "$PERSON_ATTENTION_DIM" \
+    --person-attention-heads "$PERSON_ATTENTION_HEADS" \
+    --person-history-frames "$PERSON_HISTORY_FRAMES" \
+    --person-track-max-gap-seconds "$PERSON_TRACK_MAX_GAP_SECONDS" \
+    --initialize-pooled-checkpoint "$base" \
+    --output "$checkpoint" \
+    --sessions "${SESSION_ARGS[@]}" \
+    --dev-sessions "${DEV_SESSION_ARGS[@]}" \
+    --epochs "$EPOCHS" \
+    --batch-size 512 \
+    --learning-rate 7e-5 \
+    --weight-decay 5e-5 \
+    --warmup-epochs "$WARMUP_EPOCHS" \
+    --warmup-factor 0.3 \
+    --dropout 0.2 \
+    --loss lstr-mce \
+    --transition-objective none \
+    --layer-fusion concat-mlp \
+    --expected-mode open-stream-replay \
+    --seed "$seed" \
+    --device cuda:0
+  infer_checkpoint "$checkpoint" "$result" "${SCREEN_SESSION_ARGS[@]}"
+}
+
+if [[ "$variant" == test ]]; then
+  mkdir -p "$TEST_RESULT_ROOT"
+  for selected in "${variants[@]}"; do
+    for seed in "${seeds[@]}"; do
+      checkpoint="$CHECKPOINT_ROOT/${selected}_seed${seed}.pth"
+      if [[ "$selected" == track ]]; then
+        checkpoint="$CHECKPOINT_ROOT/track_gru_seed${seed}.pth"
+        result="$TEST_RESULT_ROOT/track_gru_seed${seed}"
+      else
+        result="$TEST_RESULT_ROOT/pooled_seed${seed}"
+      fi
+      if [[ ! -s "$checkpoint" ]]; then
+        echo "Missing checkpoint: $checkpoint" >&2
+        exit 1
+      fi
+      echo "$(date): final test ${selected}, seed=${seed}"
+      infer_checkpoint "$checkpoint" "$result" 3 2 1
+    done
+  done
+  summary_root=$TEST_RESULT_ROOT
+else
+  for selected in "${variants[@]}"; do
+    for seed in "${seeds[@]}"; do
+      echo "$(date): ${selected}, seed=${seed}"
+      if [[ "$selected" == pooled ]]; then
+        train_and_infer_pooled "$seed"
+      else
+        train_and_infer_track "$seed"
+      fi
+    done
+  done
+  summary_root=$RESULT_ROOT
+fi
+
+if [[ "$SUMMARIZE" == 1 ]]; then
+  python summarize_local_baselines.py \
+    --root "$summary_root" \
+    --dataset easycom \
+    --protocol causal_valid_only \
+    --paired-deltas
+fi