vqafromscratch/run.sh at main · Dreeseaw/vqafromscratch · GitHub

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
#!/bin/bash
set -e

if [ $# -lt 1 ]; then
  echo "Usage: $0 <run_id> [checkpoint_step] [extra train args...]"
  exit 1
fi

RUN_ID="$1"
shift || true

CKPT_STEP=""
if [[ $# -gt 0 && "$1" =~ ^[0-9]+$ ]]; then
  CKPT_STEP="$1"
  shift || true
fi
EXTRA_ARGS=("$@")

mkdir -pv "logs/$RUN_ID"
if [[ -n "${CKPT_STEP}" ]]; then
  cat models/vae.py train/train.py > "logs/$RUN_ID/code_$CKPT_STEP.py"
  docker run --rm -it --gpus all --ipc=host \
    -e PYTORCH_ENABLE_MPS_FALLBACK=1 \
    -v "$(pwd)":/app -w /app myrepo:gpu \
    python -m train.train "$RUN_ID" "$CKPT_STEP" \
      "${EXTRA_ARGS[@]}"
else
  cat models/vae.py train/train.py > "logs/$RUN_ID/code.py"
  docker run --rm -it --gpus all --ipc=host \
    -e PYTORCH_ENABLE_MPS_FALLBACK=1 \
    -v "$(pwd)":/app -w /app myrepo:gpu \
    python -m train.train "$RUN_ID" \
      "${EXTRA_ARGS[@]}"
fi