File size: 940 Bytes
3ce19a2 | 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 | #!/usr/bin/env bash
set -euo pipefail
cd "$(dirname "$0")/.."
DATA_DIR="${DATA_DIR:-/path/to/AFHQ}"
SAVE_DIR="${SAVE_DIR:-./runs/afhq_rtm}"
NUM_GPUS="${NUM_GPUS:-4}"
CFG_FILE="${CFG_FILE:-src/configs/AFHQ/StyleGAN2-SPD-ADA-RTM.yaml}"
mkdir -p "$SAVE_DIR"
RESUME_ARGS=()
if [ -d "$SAVE_DIR/checkpoints" ]; then
prev_run=$(find "$SAVE_DIR/checkpoints" -mindepth 1 -maxdepth 1 -type d 2>/dev/null \
| sort | tail -n1)
if [ -n "$prev_run" ] && ls "$prev_run"/model=*.pth >/dev/null 2>&1; then
echo "[train_afhq] resuming from $prev_run"
RESUME_ARGS=(-ckpt "$prev_run")
fi
fi
CMD=(python src/main.py -t
--cfg_file "$CFG_FILE"
--data_dir "$DATA_DIR"
--save_dir "$SAVE_DIR"
-metrics fid is prdc
--num_workers 4
--print_freq 100
--save_freq 2000
--seed 0
"${RESUME_ARGS[@]}")
if [ "$NUM_GPUS" -gt 1 ]; then
"${CMD[@]}" -DDP
else
"${CMD[@]}"
fi
|