File size: 7,380 Bytes
1737ebf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | #!/bin/bash
# ──────────────────────────────────────────────────────────────────────────────
# Training script for aesthetic scoring with custom score tokens.
#
# Each image is expanded into 14 independent samples:
# - 13 per-criterion samples (score range 1–10, bins=crit_score_bins)
# - 1 total-score sample (score range 0–100, bins=total_score_bins)
#
# Assistant output format (per sample):
# <think>{reason}</think>
# The score is <scoreK> ← criteria sample
#
# The score is <totalK> ← total sample
#
# Loss = CE_loss + crit_reg_weight * Huber(criteria) + total_reg_weight * Huber(total)
# ──────────────────────────────────────────────────────────────────────────────
# export CUDA_VISIBLE_DEVICES=0
# ── Distributed configuration (auto-detect) ──────────────────────────────────
MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
MASTER_PORT=${MASTER_PORT:-$(shuf -i 20001-29999 -n 1)}
NNODES=${WORLD_SIZE:-1}
if [ -z "$NPROC_PER_NODE" ]; then
if command -v nvidia-smi &> /dev/null; then
NPROC_PER_NODE=$(nvidia-smi --list-gpus | wc -l)
else
NPROC_PER_NODE=1
fi
fi
export DS_BUILD_SHM_COMM=0 # Disable building shared memory communication ops
export DS_BUILD_OPS=0 # Disable building ops to avoid compilation errors
# ── Paths ─────────────────────────────────────────────────────────────────────
# Base model: Qwen3-VL-8B or a previously fine-tuned checkpoint
# LLM=/mnt/hdfs/zw04mlnn01/checkpoint/llm_platform/hldy/model/Qwen/Qwen3-VL-8B-Instruct/main
LLM=
# Image root directory (must contain all image_path entries in track_1_train.json)
IMAGE_DIR=zijielou/PortraitCraft/main
# Annotation file (image_path uses subdir format: images_00/xxx.jpg)
ANNOTATION=PortraitCraft/track_1_train_subdir.json
# Output
RUN_NAME="qwen3vl_aesthetic_score_tokens_0509"
OUTPUT_DIR=result/portrait/${RUN_NAME}
# DeepSpeed config
DEEPSPEED=./scripts/zero2.json
# ── Hyper-parameters ──────────────────────────────────────────────────────────
LR=2e-5
BATCH_SIZE=2 # per GPU
GRAD_ACCUM=16 # effective batch = NPROC_PER_NODE * BATCH_SIZE * GRAD_ACCUM
NUM_EPOCHS=3
# ── Score token settings ──────────────────────────────────────────────────────
CRIT_BINS=20 # <score1>…<score20> covering 1.0–10.0
TOTAL_BINS=101 # <total1>…<total101> covering 0–100
# ── Loss weights ──────────────────────────────────────────────────────────────
CE_WEIGHT=1.0
CRIT_REG_WEIGHT=0.5 # regression loss weight for criteria samples
TOTAL_REG_WEIGHT=0.5 # regression loss weight for total-score samples
REG_LOSS_TYPE=huber # huber | mse
# ── Image resolution ──────────────────────────────────────────────────────────
MAX_PIXELS=$((32 * 32 * 1024)) # ~451k pixels (~672px side)
MIN_PIXELS=$((32 * 32 * 224)) # ~12k pixels
IMAGE_MAX_TOKEN_NUM=1024
# ── Entry point ───────────────────────────────────────────────────────────────
ENTRY=qwenvl/train/train_qwen.py
# ── Update aesthetic data path in data registry ───────────────────────────────
# Patch the image_dir at runtime via environment variable so we don't hard-code
# it inside __init__.py. The AestheticScoringDataset reads IMAGE_DIR if set.
export AESTHETIC_IMAGE_DIR=${IMAGE_DIR}
export AESTHETIC_ANNOTATION=${ANNOTATION}
export WANDB_DISABLED=true
# ── Auto-compute save_steps = half epoch ──────────────────────────────────────
NUM_IMAGES=$(python -c "import json; print(len(json.load(open('${ANNOTATION}'))))")
SAMPLES_PER_IMAGE=14 # 13 criteria + 1 total
TOTAL_SAMPLES=$((NUM_IMAGES * SAMPLES_PER_IMAGE))
EFFECTIVE_BATCH=$((NPROC_PER_NODE * NNODES * BATCH_SIZE * GRAD_ACCUM))
STEPS_PER_EPOCH=$(( (TOTAL_SAMPLES + EFFECTIVE_BATCH - 1) / EFFECTIVE_BATCH ))
SAVE_STEPS=$(( STEPS_PER_EPOCH / 2 ))
[ ${SAVE_STEPS} -lt 1 ] && SAVE_STEPS=1
echo "[auto] ${NUM_IMAGES} images × ${SAMPLES_PER_IMAGE} = ${TOTAL_SAMPLES} samples"
echo "[auto] effective_batch=${EFFECTIVE_BATCH}, steps/epoch=${STEPS_PER_EPOCH}, save_steps=${SAVE_STEPS} (half epoch)"
# ── Training arguments ────────────────────────────────────────────────────────
ARGS="
--model_name_or_path ${LLM} \
--use_score_tokens True \
--crit_score_bins ${CRIT_BINS} \
--total_score_bins ${TOTAL_BINS} \
--score_ce_weight ${CE_WEIGHT} \
--crit_reg_weight ${CRIT_REG_WEIGHT} \
--total_reg_weight ${TOTAL_REG_WEIGHT} \
--tune_mm_vision True \
--tune_mm_mlp True \
--tune_mm_llm True \
--bf16 True \
--output_dir ${OUTPUT_DIR} \
--num_train_epochs ${NUM_EPOCHS} \
--per_device_train_batch_size ${BATCH_SIZE} \
--per_device_eval_batch_size $((BATCH_SIZE * 2)) \
--gradient_accumulation_steps ${GRAD_ACCUM} \
--max_pixels ${MAX_PIXELS} \
--min_pixels ${MIN_PIXELS} \
--eval_strategy no \
--save_strategy steps \
--save_steps ${SAVE_STEPS} \
--save_total_limit 6 \
--learning_rate ${LR} \
--weight_decay 0.0 \
--warmup_ratio 0.03 \
--max_grad_norm 2.0 \
--lr_scheduler_type cosine \
--logging_steps 1 \
--model_max_length 4096 \
--gradient_checkpointing False \
--dataloader_num_workers 4 \
--run_name ${RUN_NAME} \
--sample_every_steps 500 \
--report_to none \
--deepspeed ${DEEPSPEED}
"
# ── Launch ────────────────────────────────────────────────────────────────────
mkdir -p ${OUTPUT_DIR}
echo "============================================================"
echo " GPUs: ${NPROC_PER_NODE}, Nodes: ${NNODES}"
echo " Master: ${MASTER_ADDR}:${MASTER_PORT}"
echo " Output: ${OUTPUT_DIR}"
echo "============================================================"
torchrun \
--nproc_per_node=${NPROC_PER_NODE} \
--nnodes=${NNODES} \
--master_addr=${MASTER_ADDR} \
--master_port=${MASTER_PORT} \
${ENTRY} ${ARGS} 2>&1 | tee ${OUTPUT_DIR}/train.log
|