blade/test_bash_slurm.sh at main · NVlabs/blade · GitHub

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
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.


# ======================= NOTE =======================
# set ENV variables MY_CONDA_DIR, TORCH_HOME, WANDB_DIR. the later two avoids storage quota on nodes

# example usage single node:
# export NUM_NODES=1; export NUM_GPUS=8; export MINI_BATCHSIZE=2; export EXP_NAME="depth_${NUM_NODES}x${NUM_GPUS}x${MINI_BATCHSIZE}";
# torchrun --nproc_per_node $NUM_GPUS --master_addr localhost --master_port 12356 --nnodes $NUM_NODES --node_rank 0 ./scripts/train.py ./blade/configs/blade_posenet.py --launcher pytorch --work-dir=./work_dirs/train_depth_${NUM_NODES}x${NUM_GPUS}x${MINI_BATCHSIZE} > log_${NUM_NODES}x${NUM_GPUS}x${MINI_BATCHSIZE}.txt


source ~/.bashrc
: "${MY_CONDA_DIR:?Environment variable MY_CONDA_DIR must be set}"
export miniconda_folder="${MY_CONDA_DIR%/}"
export cur_repo_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
#export TORCH_HOME=/path/to/.cache/
#export WANDB_DIR=/path/to/.wandb_cache/

# >>> conda initialize >>>
# !! Contents within this block are managed by 'conda init' !!
conda_exe="$miniconda_folder/bin/conda"
__conda_setup="$('$conda_exe' 'shell.bash' 'hook' 2> /dev/null)"
if [ $? -eq 0 ]; then
    eval "$__conda_setup"
else
    if [ -f "$miniconda_folder/etc/profile.d/conda.sh" ]; then
        . "$miniconda_folder/etc/profile.d/conda.sh"
    else
        export PATH="$miniconda_folder/bin:$PATH"
    fi
fi
unset __conda_setup
# <<< conda initialize <<<
conda activate $CUR_CONDA_ENV


#export NCCL_IB_DISABLE=1
export OMP_NUM_THREADS=12
echo $CUR_CONDA_ENV
conda env list
conda list

pip install wandb mediapipe  # bug fixed: wandb need to be installed in bash script for some reason, might not be the case for other clusters

export EXP_NAME=$1
export CONFIG_FILE=$2
export WORKDIR="$cur_repo_root/work_dirs/${EXP_NAME}"
echo "SUBMIT_GPUS: ${SUBMIT_GPUS}"
echo "MASTER_ADDR: ${MASTER_ADDR}"
echo "MASTER_PORT: ${MASTER_PORT}"
echo "NUM_NODES: ${NUM_NODES}"
echo "NODE_RANK: ${NODE_RANK}"
echo "CONFIG_FILE: ${CONFIG_FILE}"
echo "WORKDIR: ${WORKDIR}"
torchrun --nproc_per_node $SUBMIT_GPUS --master_addr $MASTER_ADDR --master_port $MASTER_PORT --nnodes $NUM_NODES --node_rank $NODE_RANK ./scripts/test.py $CONFIG_FILE --launcher pytorch --work-dir="$WORKDIR" --out ${WORKDIR}/results_${TESTSET}.pkl --data-name $TESTSET --checkpoint $CKPT