run_train_official_config.sh

#!/usr/bin/env bash
# UVDoc official-style training (matches UVDoc_official/train.py defaults).
#
# Mode A — paper/repo default: Doc3D + UVDoc mixed (--data_to_use both).
#   Set DOC3D_ROOT to your Doc3D dataset; UVDoc to UVDoc_final (or official layout).
#
# Mode B — local UVDoc only: same LR/epoch/gamma schedule as official, no Doc3D.
#   export TRAIN_MODE=uvdoc_only
#
# Defaults (official argparse):
#   batch_size=8, n_epochs=10, n_epochs_decay=10, lr=2e-4,
#   alpha_w=5 beta_w=5 gamma_w=1 ep_gamma_start=10,
#   appearance: visual noise color, UVDoc geom: rotate

set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PY="${PYTHON:-python3}"
TRAIN="${ROOT}/UVDoc_official/train.py"

# --- edit paths ---
DOC3D_ROOT="${DOC3D_ROOT:-/path/to/data/doc3D}"
UV_DOC_ROOT="${UV_DOC_ROOT:-/mnt/zsn/zsn_workspace/dzx/UvDoc/UVDoc_final}"
LOGDIR="${LOGDIR:-${ROOT}/log/official_default}"

# TRAIN_MODE: both | uvdoc_only
TRAIN_MODE="${TRAIN_MODE:-both}"

# Official hyperparameters (explicit for clarity)
BS="${BS:-8}"
N_EPOCHS="${N_EPOCHS:-10}"
N_EPOCHS_DECAY="${N_EPOCHS_DECAY:-10}"
LR="${LR:-0.0002}"
ALPHA="${ALPHA:-5.0}"
BETA="${BETA:-5.0}"
GAMMA="${GAMMA:-1.0}"
EP_GAMMA_START="${EP_GAMMA_START:-10}"
UV_VAL_RATIO="${UV_VAL_RATIO:-0.05}"
UV_SPLIT_SEED="${UV_SPLIT_SEED:-42}"
NUM_WORKERS="${NUM_WORKERS:-8}"
DEVICE="${DEVICE:-cuda:0}"

common_args=(
  --batch_size "$BS"
  --n_epochs "$N_EPOCHS"
  --n_epochs_decay "$N_EPOCHS_DECAY"
  --lr "$LR"
  --alpha_w "$ALPHA"
  --beta_w "$BETA"
  --gamma_w "$GAMMA"
  --ep_gamma_start "$EP_GAMMA_START"
  --appearance_augmentation visual noise color
  --geometric_augmentationsUVDoc rotate
  --num_workers "$NUM_WORKERS"
  --device "$DEVICE"
  --logdir "$LOGDIR"
)

if [[ "$TRAIN_MODE" == "both" ]]; then
  if [[ ! -d "$DOC3D_ROOT" ]]; then
    echo "ERROR: DOC3D_ROOT is not a directory: $DOC3D_ROOT"
    echo "Set DOC3D_ROOT to your Doc3D dataset root, or use TRAIN_MODE=uvdoc_only."
    exit 1
  fi
  exec "$PY" "$TRAIN" \
    --data_to_use both \
    --data_path_doc3D "$DOC3D_ROOT" \
    --data_path_UVDoc "$UV_DOC_ROOT" \
    "${common_args[@]}"
elif [[ "$TRAIN_MODE" == "uvdoc_only" ]]; then
  exec "$PY" "$TRAIN" \
    --data_to_use uvdoc \
    --data_path_UVDoc "$UV_DOC_ROOT" \
    --uvdoc_val_ratio "$UV_VAL_RATIO" \
    --uvdoc_split_seed "$UV_SPLIT_SEED" \
    "${common_args[@]}"
else
  echo "TRAIN_MODE must be 'both' or 'uvdoc_only', got: $TRAIN_MODE"
  exit 1
fi