| #!/usr/bin/env bash |
| set -euo pipefail |
|
|
| ROOT=${ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)} |
| umask "${UCFNET_UMASK:-000}" |
|
|
| usage() { |
| cat >&2 <<EOF |
| Usage: bash scripts/train_ucfnet.sh [-g GPUS] [-c CONFIG] [-p PORT] [-r RESUME] [-b BATCH] [-n] [-h] |
| |
| -g GPUS comma-separated GPU list, e.g. 0,1,3,5 (default: 0) |
| -c CONFIG detector yaml: a bare name under |
| DeepfakeBench/training/config/detector/ (e.g. ucfnet.yaml) |
| or a path (default: ucfnet.yaml) |
| -p PORT DDP master port (default: 29617) |
| -r RESUME resume from a checkpoint (latest.pth) |
| -b BATCH train batch size per GPU |
| -n disable DDP (single process) |
| -h show this help |
| |
| Flags override the matching env vars (GPU_LIST / DETECTOR_TEMPLATE / |
| MASTER_PORT / RESUME / TRAIN_BATCH_SIZE / DDP), which still work too. |
| EOF |
| } |
|
|
| DETECTOR_DIR="$ROOT/DeepfakeBench/training/config/detector" |
| while getopts "g:c:p:r:b:nh" opt; do |
| case "$opt" in |
| g) GPU_LIST=$OPTARG ;; |
| c) if [[ "$OPTARG" == */* ]]; then DETECTOR_TEMPLATE=$OPTARG; else DETECTOR_TEMPLATE="$DETECTOR_DIR/$OPTARG"; fi ;; |
| p) MASTER_PORT=$OPTARG ;; |
| r) RESUME=$OPTARG ;; |
| b) TRAIN_BATCH_SIZE=$OPTARG ;; |
| n) DDP=0 ;; |
| h) usage; exit 0 ;; |
| *) usage; exit 1 ;; |
| esac |
| done |
| shift $((OPTIND - 1)) |
|
|
| source "$ROOT/scripts/resolve_python.sh" |
| ENV_PYTHON=$(resolve_env_python "$ROOT") |
| DATA_ROOT=${DATA_ROOT:-$ROOT/UCF-Net-dataset} |
| PRETRAINED_ROOT=${PRETRAINED_ROOT:-$ROOT/pretrained} |
| DETECTOR_TEMPLATE=${DETECTOR_TEMPLATE:-$ROOT/DeepfakeBench/training/config/detector/ucfnet.yaml} |
| TRAIN_CONFIG=${TRAIN_CONFIG:-$ROOT/DeepfakeBench/training/config/train_config.yaml} |
| RUNTIME_CONFIG=${RUNTIME_CONFIG:-$ROOT/.runtime/configs/ucfnet.yaml} |
| LOG_DIR=${LOG_DIR:-$ROOT/runs/train/ucfnet} |
| LIBFIX_DIR=${LIBFIX_DIR:-$ROOT/.libfix} |
| GPU_LIST=${GPU_LIST:-0} |
| MASTER_PORT=${MASTER_PORT:-29617} |
| DDP=${DDP:-1} |
| RESUME=${RESUME:-} |
| RESUME_WEIGHTS_ONLY=${RESUME_WEIGHTS_ONLY:-0} |
| TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-} |
| TEST_BATCH_SIZE=${TEST_BATCH_SIZE:-} |
| TRAIN_WORKERS=${TRAIN_WORKERS:-} |
| MAX_TRAIN_ITERS=${MAX_TRAIN_ITERS:-} |
|
|
| resolve_file() { |
| local path=$1 |
| if [[ "$path" = /* ]]; then |
| printf '%s\n' "$path" |
| else |
| printf '%s\n' "$ROOT/$path" |
| fi |
| } |
|
|
| resolve_dir() { |
| local path=$1 |
| if [[ "$path" = /* ]]; then |
| printf '%s\n' "$path" |
| else |
| printf '%s\n' "$ROOT/$path" |
| fi |
| } |
|
|
| DATA_ROOT=$(resolve_dir "$DATA_ROOT") |
| PRETRAINED_ROOT=$(resolve_dir "$PRETRAINED_ROOT") |
| LOG_DIR=$(resolve_dir "$LOG_DIR") |
| LIBFIX_DIR=$(resolve_dir "$LIBFIX_DIR") |
| DETECTOR_TEMPLATE=$(resolve_file "$DETECTOR_TEMPLATE") |
| TRAIN_CONFIG=$(resolve_file "$TRAIN_CONFIG") |
| RUNTIME_CONFIG=$(resolve_file "$RUNTIME_CONFIG") |
| if [[ -n "$RESUME" ]]; then |
| RESUME=$(resolve_file "$RESUME") |
| fi |
|
|
| require_env_python "$ENV_PYTHON" |
|
|
| if [[ ! -e "$LIBFIX_DIR/libcuda.so" ]]; then |
| libcuda_so1=$(ldconfig -p 2>/dev/null | awk '/libcuda\.so\.1/{print $NF; exit}') |
| if [[ -n "$libcuda_so1" ]]; then |
| mkdir -p "$LIBFIX_DIR" |
| ln -sf "$libcuda_so1" "$LIBFIX_DIR/libcuda.so" |
| fi |
| fi |
| if [[ -d "$LIBFIX_DIR" ]]; then |
| export LD_LIBRARY_PATH="$LIBFIX_DIR:/usr/lib/x86_64-linux-gnu${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" |
| fi |
|
|
| materialize_args=( |
| "$ROOT/scripts/materialize_config.py" |
| --repo-root "$ROOT" \ |
| --template "$DETECTOR_TEMPLATE" \ |
| --output "$RUNTIME_CONFIG" \ |
| --data-root "$DATA_ROOT" \ |
| --pretrained-root "$PRETRAINED_ROOT" \ |
| --log-dir "$LOG_DIR" |
| ) |
| if [[ -n "$TRAIN_BATCH_SIZE" ]]; then |
| materialize_args+=(--train-batch-size "$TRAIN_BATCH_SIZE") |
| fi |
| if [[ -n "$TEST_BATCH_SIZE" ]]; then |
| materialize_args+=(--test-batch-size "$TEST_BATCH_SIZE") |
| fi |
| if [[ -n "$TRAIN_WORKERS" ]]; then |
| materialize_args+=(--workers "$TRAIN_WORKERS") |
| fi |
| if [[ -n "$MAX_TRAIN_ITERS" ]]; then |
| materialize_args+=(--max-train-iters "$MAX_TRAIN_ITERS") |
| fi |
| "$ENV_PYTHON" "${materialize_args[@]}" >/dev/null |
|
|
| IFS=',' read -r -a GPUS <<< "$GPU_LIST" |
| NPROC=${NPROC:-${#GPUS[@]}} |
|
|
| cmd=( |
| "$ENV_PYTHON" -m torch.distributed.launch |
| --master_port="$MASTER_PORT" |
| --nproc_per_node="$NPROC" |
| training/train.py |
| --detector_path "$RUNTIME_CONFIG" |
| --train_config_path "$TRAIN_CONFIG" |
| ) |
|
|
| if [[ "$DDP" == "1" ]]; then |
| cmd+=(--ddp) |
| fi |
| if [[ -n "$RESUME" ]]; then |
| cmd+=(--resume "$RESUME") |
| fi |
| if [[ "$RESUME_WEIGHTS_ONLY" == "1" ]]; then |
| cmd+=(--resume-weights-only) |
| fi |
|
|
| mkdir -p "$LOG_DIR" |
| chmod -R a+rwX "$LOG_DIR" 2>/dev/null || true |
| cd "$ROOT/DeepfakeBench" |
| echo "Runtime config: $RUNTIME_CONFIG" |
| echo "Logs/checkpoints: $LOG_DIR" |
| CUDA_VISIBLE_DEVICES="$GPU_LIST" "${cmd[@]}" |
|
|
