train_qwen_enhanced_6gpu.sh

#!/bin/bash
# PriorCLIP-BLIP 6GPU训练脚本 - RSICD + RSITMD (使用Qwen增强文本)
# 使用 GPU 0-5 (六张GPU)

TIMESTAMP=$(date +"%m%d_%H%M")
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5  # 使用6张GPU
export PYTHONPATH="/home/yy/testPIR/src:$PYTHONPATH"
export NCCL_P2P_DISABLE=1
export NCCL_IB_DISABLE=1
export TORCH_CUDA_ARCH_LIST="9.0"

# HuggingFace 镜像设置
export HF_ENDPOINT=https://hf-mirror.com

echo "=========================================="
echo "PriorCLIP-BLIP 训练 (Qwen增强文本)"
echo "Timestamp: $TIMESTAMP"
echo "Using GPUs: 0-5 (6x GPU)"
echo "=========================================="

# 数据集配置 - 使用Qwen增强的CSV文件
# 合并后的训练集
COMBINED_TRAIN="/data/yy/combined_train_qwen_enhanced.csv"
RSICD_VAL="/data/yy/RSICD/val_qwen_enhanced.csv"

echo ""
echo "数据集配置 (Qwen增强):"
echo "  合并训练集: $COMBINED_TRAIN"
echo "  RSICD验证集: $RSICD_VAL"
echo ""

# 检查文件
echo "检查数据文件..."
ls -lh "$COMBINED_TRAIN" "$RSICD_VAL"
echo ""

# 统计样本数
TOTAL_TRAIN_SAMPLES=$(($(wc -l < "$COMBINED_TRAIN") - 1))
RSICD_VAL_SAMPLES=$(($(wc -l < "$RSICD_VAL") - 1))

echo "数据集统计:"
echo "  总训练样本: $TOTAL_TRAIN_SAMPLES"
echo "  RSICD验证样本: $RSICD_VAL_SAMPLES"
echo ""

echo "Starting 6-GPU training..."

# 创建日志目录
LOG_DIR="./logs/QwenEnhanced_RSICD_RSITMD_6GPU_${TIMESTAMP}"
mkdir -p "$LOG_DIR"

# 启动6 GPU训练
# 切换到包含 src 目录的父目录，以便 python -m 可以找到 training.main
cd /home/yy/testPIR

conda run -n testPIR --no-capture-output python -m torch.distributed.run \
    --nproc_per_node=6 \
    --master_port=29500 \
    -m training.main -- \
    --model PriorCLIP_BLIP \
    --train-data "$COMBINED_TRAIN" \
    --val-data "$RSICD_VAL" \
    --dataset-type csv \
    --csv-img-key image_path \
    --csv-caption-key caption \
    --csv-separator "," \
    --train-num-samples $TOTAL_TRAIN_SAMPLES \
    --val-num-samples $RSICD_VAL_SAMPLES \
    --batch-size 32 \
    --accum-freq 2 \
    --epochs 50 \
    --lr 5e-4 \
    --beta1 0.9 \
    --beta2 0.98 \
    --eps 1e-6 \
    --warmup 500 \
    --wd 0.2 \
    --precision amp_bf16 \
    --local-loss \
    --gather-with-grad \
    --save-frequency 1 \
    --save-most-recent \
    --val-frequency 1 \
    --zeroshot-frequency 1 \
    --log-every-n-steps 10 \
    --seed 0 \
    --aug-cfg scale='(0.4, 1.0)' color_jitter='(0.32, 0.32, 0.32, 0.08)' color_jitter_prob=0.8 gray_scale_prob=0.2 \
    --name "QwenEnhanced_RSICD_RSITMD_6GPU_${TIMESTAMP}" \
    --report-to tensorboard \
    --logs ./logs \
    2>&1 | tee "$LOG_DIR/train_output.log"

echo ""
echo "=========================================="
echo "训练已启动！"
echo "日志目录: $LOG_DIR/"
echo ""
echo "监控训练:"
echo "  tensorboard --logdir ./logs/QwenEnhanced_RSICD_RSITMD_6GPU_${TIMESTAMP}/"
echo "  tail -f $LOG_DIR/train_output.log"
echo ""
echo "检查GPU使用:"
echo "  watch -n 1 nvidia-smi"
echo "=========================================="