VLAlert / training /PRETRAIN /run_pretrain.sh
AsianPlayer's picture
Add VLAlert code
1e05592 verified
Raw
History Blame Contribute Delete
1.42 kB
#!/bin/bash
# VLM预训练启动脚本
# 设置环境变量
export CUDA_VISIBLE_DEVICES=0
export PYTHONPATH="PROJECT_ROOT:$PYTHONPATH"
# 训练目录
TRAIN_DIR="PROJECT_ROOT/training/pretrain"
mkdir -p $TRAIN_DIR
cd $TRAIN_DIR
echo "======================================"
echo "VLM预训练"
echo "======================================"
echo ""
# 检查数据
DATA_FILE="PROJECT_ROOT/data/dataset/pretrain/train/pretrain_data.pkl"
if [ ! -f "$DATA_FILE" ]; then
echo "❌ 数据文件不存在: $DATA_FILE"
echo "请先运行: python prepare_pretrain_data.py"
exit 1
fi
echo "✓ 数据文件: $DATA_FILE"
# 创建输出目录
OUTPUT_DIR="PROJECT_ROOT/checkpoints/pretrain"
mkdir -p $OUTPUT_DIR
echo "✓ 输出目录: $OUTPUT_DIR"
echo ""
# 选择模型
MODEL=$1
if [ -z "$MODEL" ]; then
echo "用法: bash run_pretrain.sh [qwen2.5-vl-3b|qwen2.5-vl-7b]"
echo ""
echo "示例:"
echo " bash run_pretrain.sh qwen2.5-vl-3b"
echo " bash run_pretrain.sh qwen2.5-vl-7b"
exit 1
fi
echo "======================================"
echo "开始训练: $MODEL"
echo "======================================"
echo ""
# 训练
python train_pretrain.py \
--model $MODEL \
--epochs 5 \
--batch_size 1 \
--lr 2e-5
echo ""
echo "======================================"
echo "训练完成!"
echo "======================================"
echo "Checkpoints: $OUTPUT_DIR/$MODEL"