run2.sh

#!/bin/bash

a=$(echo $HOSTNAME | cut  -c12-16)

CONFIG=$1
JOB_NAME=${2:-"experiments"}
GPUS=${3:-8}
  
partition=${4:-'local'} #

GPUS_PER_NODE=${GPUS:-8}
if [ $GPUS_PER_NODE -ge 8 ]; then
  GPUS_PER_NODE=8
fi
CPUS_PER_TASK=${CPUS_PER_TASK:-4}
SRUN_ARGS=${SRUN_ARGS:-""}

PY_ARGS=${@:5}

WORK_DIR=${CONFIG//configs/work_dirs}
WORK_DIR=${WORK_DIR//.yaml//$JOB_NAME}
echo $WORK_DIR
mkdir  -p $WORK_DIR
mkdir -p data/temp

# please change DATA_PATH where you put the training data
export DATA_PATH='/root/autodl-tmp/datas'

# 无 SLURM 时，使用 torchrun 启动；即使 1 卡也会初始化默认进程组，避免采样器报错
torchrun --standalone --nproc_per_node=${GPUS} main.py \
--num-gpus $GPUS \
--config-file ${CONFIG} --init_method pytorch --resume \
${PY_ARGS} OUTPUT_DIR $WORK_DIR