File size: 2,303 Bytes
f693366 | 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 | #!/bin/bash
# Atlas online training: detection + planning + caption.
# Default: visual_token_mode=online, live frozen StreamPETR + TopoMLP.
# Data: 28k det + 24k plan + 169k caption = ~220k samples/epoch
# WARNING: caption data is 6x larger than other tasks — consider downsampling.
#
# Usage:
# bash scripts/train_with_caption_balanced.sh
# RESUME_CKPT=work_dirs/atlas_with_caption/epoch-4/checkpoint.pt bash scripts/train_with_caption_balanced.sh
# NUM_GPUS=8 bash scripts/train_with_caption_balanced.sh
set -e
PROJECT_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
cd "$PROJECT_ROOT"
ENV_ROOT="${PROJECT_ROOT}/envs/streampetr"
DEEPSPEED_BIN="${ENV_ROOT}/bin/deepspeed"
if [ ! -x "$DEEPSPEED_BIN" ]; then
echo "ERROR: deepspeed not found in ${DEEPSPEED_BIN}" >&2
exit 1
fi
export PATH="${ENV_ROOT}/bin:${PATH}"
export LD_LIBRARY_PATH="${ENV_ROOT}/lib:${LD_LIBRARY_PATH}"
NUM_GPUS=${NUM_GPUS:-4}
PLANNING_TABLE3_MODE=${PLANNING_TABLE3_MODE:-atlas_high_level_ego}
TASK_LOSS_WEIGHTS=${TASK_LOSS_WEIGHTS:-detection=0.35,planning=1.0,caption=0.05}
OUTPUT_DIR=${OUTPUT_DIR:-work_dirs/atlas_with_caption_online}
EXTRA_ARGS=""
if [ -n "$RESUME_CKPT" ]; then
EXTRA_ARGS="--resume $RESUME_CKPT"
fi
"$DEEPSPEED_BIN" --num_gpus "$NUM_GPUS" train_atlas.py \
--llm_model pretrained/vicuna-7b-v1.5 \
--data_json data/atlas_nuscenes_train.json,data/atlas_planning_train_uniad_command.json,data/atlas_caption_train_canonical.json \
--data_root /home/guoyuanbo/autodl-tmp/data/nuscenes \
--visual_token_mode online \
--task_balance_mode scene_unit_111 \
--task_loss_weights "$TASK_LOSS_WEIGHTS" \
--planning_table3_mode "$PLANNING_TABLE3_MODE" \
--streampetr_config configs/streampetr_atlas_aligned.py \
--streampetr_ckpt pretrained/streampetr/streampetr_eva02_ep24.pth \
--topomlp_config configs/topomlp_atlas_aligned.py \
--topomlp_ckpt work_dirs/topomlp_atlas_aligned/epoch_24.pth \
--deepspeed configs/ds_zero2.json \
--output_dir "$OUTPUT_DIR" \
--epochs 8 \
--lr 2e-5 \
--weight_decay 1e-4 \
--batch_size 1 \
--gradient_accumulation_steps 2 \
--warmup_ratio 0.03 \
--max_grad_norm 1.0 \
--log_steps 100 \
--save_epochs 1 \
--keep_last_n_ckpts 3 \
--seed 42 \
--num_workers 4 \
$EXTRA_ARGS
|