lora_med_sapiens.sh

#!/bin/bash
cd src/pose

###--------------------------------------------------------------
# DEVICES and port
DEVICES=0,1,2,3,
#DEVICES=0,1,

RUN_FILE='./tools/dist_train.sh'
PORT=$(( ((RANDOM<<15)|RANDOM) % 63001 + 2000 ))

###--------------------------------------------------------------
## PARAMETERS
DATASET=$1  # Pass the dataset as the first argument to the script
MODEL="lora_med_sapiens_0.3b-210e_${DATASET}-1024x768"
JOB_NAME="pose_$MODEL"
TRAIN_BATCH_SIZE_PER_GPU=2

RESUME_FROM=''
LOAD_FROM=''

# Debugging mode or multi-GPU mode
mode=${2:-'multi-gpu'}  # Default mode is multi-gpu if not specified

###--------------------------------------------------------------
CONFIG_FILE=configs/lora_med_sapiens/${DATASET}/${MODEL}.py
OUTPUT_DIR="Outputs/train_lora/${DATASET}/${MODEL}/node"
OUTPUT_DIR="$(echo "${OUTPUT_DIR}/$(date +"%m-%d-%Y_%H:%M:%S")")"

###--------------------------------------------------------------
if [ -n "$LOAD_FROM" ]; then
    OPTIONS="train_dataloader.batch_size=$TRAIN_BATCH_SIZE_PER_GPU load_from=$LOAD_FROM"
else
    OPTIONS="train_dataloader.batch_size=$TRAIN_BATCH_SIZE_PER_GPU"
fi

if [ -n "$RESUME_FROM" ]; then
    CMD_RESUME="--resume ${RESUME_FROM}"
else
    CMD_RESUME=""
fi

export TF_CPP_MIN_LOG_LEVEL=2

###--------------------------------------------------------------
## Run the appropriate mode
if [ "$mode" = "debug" ]; then
    TRAIN_BATCH_SIZE_PER_GPU=2
    OPTIONS="$(echo "train_dataloader.batch_size=${TRAIN_BATCH_SIZE_PER_GPU} train_dataloader.num_workers=0 train_dataloader.persistent_workers=False")"

    CUDA_VISIBLE_DEVICES=${DEVICES} python tools/train.py ${CONFIG_FILE} \
        --work-dir ${OUTPUT_DIR} \
        --no-validate \
        --cfg-options ${OPTIONS}

elif [ "$mode" = "multi-gpu" ]; then
    NUM_GPUS_STRING_LEN=${#DEVICES}
    NUM_GPUS=$((NUM_GPUS_STRING_LEN / 2))
    SEED='0'

    LOG_FILE="$(echo "${OUTPUT_DIR}/log.txt")"
    mkdir -p ${OUTPUT_DIR}; touch ${LOG_FILE}

    CUDA_VISIBLE_DEVICES=${DEVICES} PORT=${PORT} ${RUN_FILE} ${CONFIG_FILE} \
        ${NUM_GPUS} \
        --work-dir ${OUTPUT_DIR} \
        --seed ${SEED} \
        --cfg-options ${OPTIONS} \
        ${CMD_RESUME} \
        | tee ${LOG_FILE}
fi