File size: 1,546 Bytes
3ce19a2 59ee8cd 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 40 41 42 43 44 45 | #!/usr/bin/env bash
set -euo pipefail
DATA_ROOT="${DATA_ROOT:-./datasets/cifar10}"
SAVE_DIR="${SAVE_DIR:-./runs/cifar10_rtm}"
NUM_GPUS="${NUM_GPUS:-1}"
mkdir -p "$SAVE_DIR"
CKPT_DIR="$SAVE_DIR/train"
RESTORE_ARGS=()
if [ -f "$CKPT_DIR/latest-model.th" ]; then
echo "[train_cifar10] resuming from $CKPT_DIR/latest-model.th"
RESTORE_ARGS+=(--restore_path "$CKPT_DIR/latest-model.th")
[ -f "$CKPT_DIR/latest-model-ema.th" ] && RESTORE_ARGS+=(--restore_ema_path "$CKPT_DIR/latest-model-ema.th")
[ -f "$CKPT_DIR/latest-opt.th" ] && RESTORE_ARGS+=(--restore_optimizer_path "$CKPT_DIR/latest-opt.th")
[ -f "$CKPT_DIR/latest-sched.th" ] && RESTORE_ARGS+=(--restore_scheduler_path "$CKPT_DIR/latest-sched.th")
[ -f "$CKPT_DIR/latest-log.jsonl" ] && RESTORE_ARGS+=(--restore_log_path "$CKPT_DIR/latest-log.jsonl")
fi
torchrun --nnodes=1 --nproc_per_node="$NUM_GPUS" --standalone train.py \
--hps cifar10 \
--save_dir "$SAVE_DIR" \
--data_root "$DATA_ROOT" \
--num_epochs 200 \
--fid_freq 1000 \
--use_se True \
--width 256 \
--dec_blocks '1x1,4m1,4x2,8m4,8x2,16m8,16x2,32m16,32x2' \
--force_factor 5 \
--imle_force_resample 5 \
--lr 0.0008 \
--search_type lpips \
--n_batch 256 \
--imle_batch 1024 \
--iters_per_save 1000 \
--iters_per_images 5000 \
--iters_per_ckpt 100000 \
--latent_dim 128 \
--use_rtm True \
--H_cycles 4 --L_cycles 1 --refinement_steps 4 \
--num_tokens 4 \
--rtm_hidden_size 128 \
"${RESTORE_ARGS[@]}"
|