| #!/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[@]}" |
|
|
