scripts/vm_setup.sh - master - kimudev/rubicon_test

scripts/vm_setup.sh
91 строка4 KB
Rubicon
Распознавание людей в ИК-диапазоне с беспилотника
04 авг 2026, 08:33
04 авг 2026, 08:33e8c6dfe
#!/usr/bin/env bash
#
# Разворачивает окружение на созданной виртуальной машине: монтирует диск
# с данными, ставит PyTorch и проект, скачивает оба датасета из публичных
# источников и готовит YOLO-выборки.
#
# Использование:
#   ./scripts/vm_setup.sh <IP> [путь_к_приватному_ключу]
#
set -euo pipefail

IP="${1:?укажите адрес машины: ./scripts/vm_setup.sh 1.2.3.4}"
KEY="${2:-$HOME/.ssh/id_ed25519}"
SSH=(ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -i "$KEY" "yc-user@$IP")

echo "==> проверяю доступ"
"${SSH[@]}" 'nvidia-smi --query-gpu=name,memory.total --format=csv,noheader'

echo "==> копирую проект"
tar czf /tmp/rubicon_repo.tgz --exclude=.git --exclude=.venv --exclude=artifacts .
scp -q -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -i "$KEY" \
    /tmp/rubicon_repo.tgz "yc-user@$IP:~/"

"${SSH[@]}" 'bash -s' <<'REMOTE'
set -euo pipefail

echo "==> монтирую диск с данными"
if ! mountpoint -q /mnt/data; then
  DEV=$(lsblk -dpno NAME,TYPE | awk '$2=="disk"{print $1}' | tail -1)
  sudo blkid "$DEV" >/dev/null 2>&1 || sudo mkfs.ext4 -q -F "$DEV"
  sudo mkdir -p /mnt/data && sudo mount "$DEV" /mnt/data
  UUID=$(sudo blkid -s UUID -o value "$DEV")
  grep -q "$UUID" /etc/fstab || \
    echo "UUID=$UUID /mnt/data ext4 defaults,nofail 0 2" | sudo tee -a /etc/fstab >/dev/null
  sudo chown -R "$USER:$USER" /mnt/data
fi
mkdir -p /mnt/data/{archives,raw,dataset,runs}

echo "==> ставлю системные пакеты и окружение"
sudo apt-get -qq update
sudo DEBIAN_FRONTEND=noninteractive apt-get -qq install -y python3-venv unzip wget >/dev/null
[ -d ~/venv ] || python3 -m venv ~/venv
~/venv/bin/pip install -q --upgrade pip
~/venv/bin/pip install -q torch torchvision --index-url https://download.pytorch.org/whl/cu121
mkdir -p ~/rubicon && tar xzf ~/rubicon_repo.tgz -C ~/rubicon
~/venv/bin/pip install -q -e ~/rubicon
~/venv/bin/pip install -q gdown
mkdir -p ~/.config/Ultralytics          # иначе ultralytics ругается на права

# Конфиг для машины держим отдельно: путями он отличается от локального,
# а повторное копирование проекта не должно его затирать.
sed -e 's|^  raw: .*|  raw: /mnt/data/raw/RGBTDronePerson|' \
    -e 's|^  dataset: .*|  dataset: /mnt/data/dataset|' \
    -e 's|^  runs: .*|  runs: /mnt/data/runs|' \
    ~/rubicon/configs/default.yaml > /mnt/data/vm.yaml

echo "==> скачиваю датасеты (публичные источники, авторизация не нужна)"
cd /mnt/data/archives

# RGBTDronePerson — раздача авторов, Zhang et al., ISPRS 2023.
# В папке пять файлов, поэтому ограничение gdown на 50 не мешает.
if [ ! -d /mnt/data/raw/RGBTDronePerson ]; then
  ~/venv/bin/gdown --folder \
    https://drive.google.com/drive/folders/1Mi3NXQ-YG1iiIWkPbe3GQoDK68dARMN6 \
    -O /mnt/data/archives/rgbt
  unzip -q -o /mnt/data/archives/rgbt/RGBTDronePerson.zip -d /mnt/data/raw/
  mkdir -p /mnt/data/raw/RGBTDronePerson
  mv /mnt/data/raw/train /mnt/data/raw/val /mnt/data/raw/RGBTDronePerson/
  rm -rf /mnt/data/raw/__MACOSX
  find /mnt/data/raw -name '.DS_Store' -delete
fi

# HIT-UAV — для проверки на постороннем датасете, релиз на GitHub.
if [ ! -d /mnt/data/raw/hituav ]; then
  wget -q --show-progress \
    https://github.com/suojiashun/HIT-UAV-Infrared-Thermal-Dataset/releases/download/v1.2.1/HIT-UAV.zip \
    -O HIT-UAV.zip
  unzip -q -o HIT-UAV.zip -d /mnt/data/raw/hituav
fi

echo "==> готовлю выборки"
export YOLO_CONFIG_DIR=$HOME/.config/Ultralytics
cd ~/rubicon
~/venv/bin/rubicon -c /mnt/data/vm.yaml budget
~/venv/bin/rubicon -c /mnt/data/vm.yaml prepare

echo
echo "Готово. Обучение запускать так:"
echo "  cd ~/rubicon && ~/venv/bin/rubicon -c /mnt/data/vm.yaml train"
echo "Длинные задачи — через setsid nohup, иначе обрыв SSH их убьёт."
REMOTE