run_hpc.sh

#!/bin/bash
# G2VTCR v2 Training on HPC (gw3: 2x RTX 4090)
# Uses only 1 GPU, limits to 128GB RAM
set -e

cd ~/Researches/G2VTCR
source .venv/bin/activate

export CUDA_VISIBLE_DEVICES=0  # Use only 1 GPU (shared machine)

echo "=========================================="
echo " G2VTCR v2 Training Pipeline"
echo " Host: $(hostname)"
echo " GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader -i 0)"
echo " Date: $(date)"
echo "=========================================="

MODE=${1:-baseline}  # baseline or augmented

if [ "$MODE" = "augmented" ]; then
    echo " Mode: AUGMENTED (crystal + TCR-FOLD predictions)"
    echo "=========================================="
    # Phase A: Retrain contact predictor with augmented data, then benchmark
    python -u run_train.py \
        --benchmarks tchard,immrep23,epytope \
        --tcr-fold-dir ~/tcr-fold \
        --force-retrain \
        --pretrain-epochs 200 \
        --train-epochs 50 \
        --hidden-dim 128 \
        --batch-size 32 \
        --device cuda \
        --output-dir checkpoints \
        2>&1 | tee checkpoints/train_augmented_$(date +%Y%m%d_%H%M%S).txt
else
    echo " Mode: BASELINE (crystal structures only)"
    echo "=========================================="
    # Original: skip pretrain (use saved checkpoint), run benchmarks
    python -u run_train.py \
        --benchmarks tchard,immrep23,epytope \
        --skip-pretrain \
        --train-epochs 50 \
        --hidden-dim 128 \
        --batch-size 32 \
        --device cuda \
        --output-dir checkpoints \
        2>&1 | tee checkpoints/train_log_$(date +%Y%m%d_%H%M%S).txt
fi

echo "Training complete at $(date)"