# ------------------------------------------------------------------------ | |
# Copyright (c) 2022 megvii-research. All Rights Reserved. | |
# ------------------------------------------------------------------------ | |
# Modified from Deformable DETR (https://github.com/fundamentalvision/Deformable-DETR) | |
# Copyright (c) 2020 SenseTime. All Rights Reserved. | |
# ------------------------------------------------------------------------ | |
# Modified from DETR (https://github.com/facebookresearch/detr) | |
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved | |
# ------------------------------------------------------------------------ | |
set -x | |
PARTITION=$1 | |
JOB_NAME=$2 | |
GPUS=$3 | |
RUN_COMMAND=${@:4} | |
if [ $GPUS -lt 8 ]; then | |
GPUS_PER_NODE=${GPUS_PER_NODE:-$GPUS} | |
else | |
GPUS_PER_NODE=${GPUS_PER_NODE:-8} | |
fi | |
CPUS_PER_TASK=${CPUS_PER_TASK:-4} | |
SRUN_ARGS=${SRUN_ARGS:-""} | |
srun -p ${PARTITION} \ | |
--job-name=${JOB_NAME} \ | |
--gres=gpu:${GPUS_PER_NODE} \ | |
--ntasks=${GPUS} \ | |
--ntasks-per-node=${GPUS_PER_NODE} \ | |
--cpus-per-task=${CPUS_PER_TASK} \ | |
--kill-on-bad-exit=1 \ | |
${SRUN_ARGS} \ | |
${RUN_COMMAND} | |