ra2 / run.sh

Ra sdxs-VAE (step 61000)

13c76be verified 3 days ago

3.18 kB

#!/bin/bash
# Обучение на TFRecord-шардах i1 — одной командой.
#
#   bash run.sh /tmp/hf/tst 10                       # 10 эпох по всем шардам в папке
#   bash run.sh /tmp/hf/tst 10 --init /path/to/step_061000.pt
#   bash run.sh /tmp/hf/tst 10 --step 61000 --batch 8 --fg
#
# Аргументы можно передавать и флагами, и через env (флаг приоритетнее):
#   --init PATH   стартовые веса: папка transformer/ или .pt-чекпойнт  (env INIT)
#   --step N      стартовый шаг счётчика и расписания lr               (env STEP)
#   --batch N     батч                                                  (env BATCH)
#   --out DIR     куда писать чекпойнты/лог                             (env OUT)
#   --epochs N    бюджет в эпохах (второй позиционный аргумент)         (env EPOCHS)
#   --fg          работать в переднем плане
# Для .pt-чекпойнта шаг берётся из файла, --step не нужен.
set -euo pipefail
cd "$(dirname "$0")"

DATA=""; EPOCHS="${EPOCHS:-10}"
INIT="${INIT:-transformer}"; STEP="${STEP:-61000}"
OUT="${OUT:-tmp/run}"; BATCH="${BATCH:-16}"; LOG="${LOG:-train.log}"; FG=0

while [ $# -gt 0 ]; do
  case "$1" in
    --init)   INIT="$2"; shift 2 ;;
    --step)   STEP="$2"; shift 2 ;;
    --batch)  BATCH="$2"; shift 2 ;;
    --out)    OUT="$2"; shift 2 ;;
    --epochs) EPOCHS="$2"; shift 2 ;;
    --log)    LOG="$2"; shift 2 ;;
    --fg)     FG=1; shift ;;
    -*)       echo "неизвестный флаг: $1"; exit 1 ;;
    *)        if [ -z "$DATA" ]; then DATA="$1"; else EPOCHS="$1"; fi; shift ;;
  esac
done

[ -n "$DATA" ] || { echo "укажи папку с TFRecord-шардами, например: bash run.sh /tmp/hf/tst 10"; exit 1; }
[ -e "$INIT" ] || { echo "нет стартовых весов: $INIT"; exit 1; }
[ -d "$DATA" ] || { echo "нет папки с шардами: $DATA"; exit 1; }
mkdir -p "$OUT"

STEP_ARGS=(--step "$STEP")
if [ -f "$INIT" ]; then                     # .pt помнит свой шаг — не перебиваем
  STEP_ARGS=()
  echo "стартовые веса:   $INIT (.pt, шаг из файла)"
else
  echo "стартовые веса:   $INIT (шаг $STEP)"
fi
echo "shard'ов в папке: $(ls -1 "$DATA"/*.tfrecord-* 2>/dev/null | wc -l)"
echo "бюджет:           $EPOCHS эпох, батч $BATCH, out $OUT"
echo "свободно на диске: $(df -h "$OUT" | awk 'NR==2 {print $4}')"

CMD=(python3 src/train.py
  --init "$INIT" "${STEP_ARGS[@]}"
  --tfrecord "$DATA" --epochs "$EPOCHS"
  --batch "$BATCH" --size 256 --text-len 192 --prefetch 256 --workers 8
  --lr 1e-4 --lr-min 1e-5 --lr-warmup 0 --ema 0.999
  --sample-every 2000 --save-every 4000 --log-every 200
  --out "$OUT")

echo "команда:          ${CMD[*]}"
echo
if [ "$FG" = "1" ]; then
  "${CMD[@]}"
else
  setsid nohup "${CMD[@]}" < /dev/null > "$LOG" 2>&1 &
  sleep 20
  echo "запущено в фоне. смотреть:  tail -f $LOG"
fi

	#!/bin/bash
	# Обучение на TFRecord-шардах i1 — одной командой.
	#
	# bash run.sh /tmp/hf/tst 10 # 10 эпох по всем шардам в папке
	# bash run.sh /tmp/hf/tst 10 --init /path/to/step_061000.pt
	# bash run.sh /tmp/hf/tst 10 --step 61000 --batch 8 --fg
	#
	# Аргументы можно передавать и флагами, и через env (флаг приоритетнее):
	# --init PATH стартовые веса: папка transformer/ или .pt-чекпойнт (env INIT)
	# --step N стартовый шаг счётчика и расписания lr (env STEP)
	# --batch N батч (env BATCH)
	# --out DIR куда писать чекпойнты/лог (env OUT)
	# --epochs N бюджет в эпохах (второй позиционный аргумент) (env EPOCHS)
	# --fg работать в переднем плане
	# Для .pt-чекпойнта шаг берётся из файла, --step не нужен.
	set -euo pipefail
	cd "$(dirname "$0")"

	DATA=""; EPOCHS="${EPOCHS:-10}"
	INIT="${INIT:-transformer}"; STEP="${STEP:-61000}"
	OUT="${OUT:-tmp/run}"; BATCH="${BATCH:-16}"; LOG="${LOG:-train.log}"; FG=0

	while [ $# -gt 0 ]; do
	case "$1" in
	--init) INIT="$2"; shift 2 ;;
	--step) STEP="$2"; shift 2 ;;
	--batch) BATCH="$2"; shift 2 ;;
	--out) OUT="$2"; shift 2 ;;
	--epochs) EPOCHS="$2"; shift 2 ;;
	--log) LOG="$2"; shift 2 ;;
	--fg) FG=1; shift ;;
	-*) echo "неизвестный флаг: $1"; exit 1 ;;
	*) if [ -z "$DATA" ]; then DATA="$1"; else EPOCHS="$1"; fi; shift ;;
	esac
	done

	[ -n "$DATA" ] \|\| { echo "укажи папку с TFRecord-шардами, например: bash run.sh /tmp/hf/tst 10"; exit 1; }
	[ -e "$INIT" ] \|\| { echo "нет стартовых весов: $INIT"; exit 1; }
	[ -d "$DATA" ] \|\| { echo "нет папки с шардами: $DATA"; exit 1; }
	mkdir -p "$OUT"

	STEP_ARGS=(--step "$STEP")
	if [ -f "$INIT" ]; then # .pt помнит свой шаг — не перебиваем
	STEP_ARGS=()
	echo "стартовые веса: $INIT (.pt, шаг из файла)"
	else
	echo "стартовые веса: $INIT (шаг $STEP)"
	fi
	echo "shard'ов в папке: $(ls -1 "$DATA"/.tfrecord- 2>/dev/null \| wc -l)"
	echo "бюджет: $EPOCHS эпох, батч $BATCH, out $OUT"
	echo "свободно на диске: $(df -h "$OUT" \| awk 'NR==2 {print $4}')"

	CMD=(python3 src/train.py
	--init "$INIT" "${STEP_ARGS[@]}"
	--tfrecord "$DATA" --epochs "$EPOCHS"
	--batch "$BATCH" --size 256 --text-len 192 --prefetch 256 --workers 8
	--lr 1e-4 --lr-min 1e-5 --lr-warmup 0 --ema 0.999
	--sample-every 2000 --save-every 4000 --log-every 200
	--out "$OUT")

	echo "команда: ${CMD[*]}"
	echo
	if [ "$FG" = "1" ]; then
	"${CMD[@]}"
	else
	setsid nohup "${CMD[@]}" < /dev/null > "$LOG" 2>&1 &
	sleep 20
	echo "запущено в фоне. смотреть: tail -f $LOG"
	fi