Fix benchmarking scripts (#1005)

* launch_eval_blender.sh: fix script

- add shebang
- add -s option to launch a single job per GPU
- last GPU was ignored
- kill all subprocesses when script is terminated

* launch_train_blender.sh: fix script

- add -s option to launch a single job per GPU
- add -v option to use tensorboard instead of wandb
- set nerfacto options according to https://github.com/nerfstudio-project/nerfstudio/issues/806#issuecomment-1284327844
- use a single timestamp for all training jobs
- last GPU was ignored
- kill all subprocesses when script is terminated
- print the eval script command-line

* launch_eval_blender.sh: fix script

* Update benchmarking.md

* launch_train_blender.sh: add -s to eval command-line

* Update launch_train_blender.sh

* Update benchmarking.md

docs/developer_guides/debugging_tools/benchmarking.md

@@ -9,13 +9,15 @@ To launch training jobs automatically on each of these items, you can call:
 
 ```bash
 
-./scripts/benchmarking/launch_train_blender.sh -m {METHOD_NAME} -g {GPU_LIST}
+./scripts/benchmarking/launch_train_blender.sh -m {METHOD_NAME} [-s] [-v {VIS}] [{GPU_LIST}]
 ```
 
 Simply replace the arguments in brackets with the correct arguments.
 
-- `-m`: name of the method you want to benchmark (e.g. `nerfacto`, `mipnerf`).
-- `-g`: specify the list of gpus you want to use on your machine space separated. for instance, if you want to use GPU's 0-3, you will need to pass in `0 1 2 3`. If left empty, the script will automatically find available GPU's and distribute training jobs on the available GPUs.
+- `-m {METHOD_NAME}`: Name of the method you want to benchmark (e.g. `nerfacto`, `mipnerf`).
+- `-s`: Launch a single job per GPU.
+- `-v {VIS}`: Use another visualization than wandb, which is the default. Only other option is tensorboard.
+- `{GPU_LIST}`: (optional) Specify the list of gpus you want to use on your machine space separated. for instance, if you want to use GPU's 0-3, you will need to pass in `0 1 2 3`. If left empty, the script will automatically find available GPU's and distribute training jobs on the available GPUs.
   
 :::{admonition} Tip
 :class: info
@@ -26,12 +28,12 @@ To view all the arguments and annotations, you can run `./scripts/benchmarking/l
 
 A full example would be:
 
-- With `-g` specifying gpus
+- Specifying gpus
     ```bash
-    ./scripts/benchmarking/launch_train_blender.sh -m nerfacto -g 0 1 2 3
+    ./scripts/benchmarking/launch_train_blender.sh -m nerfacto 0 1 2 3
     ```
 
-- Without `-g` automatically finds available gpus
+- Automatically find available gpus
     ```bash
     ./scripts/benchmarking/launch_train_blender.sh -m nerfacto
     ```
@@ -63,7 +65,7 @@ If we wanted to run the benchmark on all the blender data for the above example,
 
 ```bash
 
-./scripts/benchmarking/launch_eval_blender.sh -m instant-ngp -o outputs/ -t 2022-08-10_172517   
+./scripts/benchmarking/launch_eval_blender.sh -m instant-ngp -o outputs/ -t 2022-08-10_172517 [{GPU_LIST}]
 ```
 
 The flags used in the benchmarking script are defined as follows:
@@ -71,7 +73,8 @@ The flags used in the benchmarking script are defined as follows:
 - `-m`: config name (e.g. `instant-ngp`). This should be the same as what was passed in for -c in the train script.
 - `-o`: base output directory for where all of the benchmarks are stored (e.g. `outputs/`). Corresponds to the `--output-dir` in the base `Config` for training.
 - `-t`: timestamp of benchmark; also the identifier (e.g. `2022-08-10_172517`).
-- `-g`: specifies the gpus to use and if not specified (no -g flag), will automaticaly search for available gpus.
+- `-s`: Launch a single job per GPU.
+- `{GPU_LIST}`: (optional) Specify the list of gpus you want to use on your machine space separated. For instance, if you want to use GPU's 0-3, you will need to pass in `0 1 2 3`. If left empty, the script will automatically find available GPU's and distribute evaluation jobs on the available GPUs.
 
 The script will simultaneously run the benchmarking across all the objects in the blender dataset and calculates the PSNR/FPS/other stats. The results are saved as .json files in the `-o` directory with the following format:
 
@@ -87,4 +90,4 @@ outputs
 :class: warning
 
 Since we are running multiple backgrounded processes concurrently with this script, please note the terminal logs may be messy.
-  :::
+  :::

scripts/benchmarking/launch_eval_blender.sh

@@ -1,19 +1,23 @@
+#!/bin/bash
+
 helpFunction_launch_eval()
 {
-   echo "Usage: $0 -m method_name -o output_dir -t timestamp -g [OPTIONAL] gpu_list"
+   echo "Usage: $0 -m <method_name> -o <output_dir> -t <timestamp> [-s] [<gpu_list>]"
    echo -e "\t-m name of method to benchmark (e.g. nerfacto, instant-ngp)"
    echo -e "\t-o base directory for where all the benchmarks are stored (e.g. outputs/)"
-   echo -e "\t-t timstamp, if using launch_train_blender.sh will be of format %Y-%m-%d_%H%M%S"
-   echo -e "\t-g [OPTIONAL] list of space-separated gpu numbers to launch train on (e.g. 0 2 4 5)"
+   echo -e "\t-t <timestamp>: if using launch_train_blender.sh will be of format %Y-%m-%d_%H%M%S"
+   echo -e "\t-s: Launch a single evaluation job per gpu."
+   echo -e "\t<gpu_list> [OPTIONAL] list of space-separated gpu numbers to launch train on (e.g. 0 2 4 5)"
    exit 1 # Exit program after printing help
 }
 
-while getopts "m:o:t:g" opt; do
+single=false
+while getopts "m:o:t:s" opt; do
     case "$opt" in
         m ) method_name="$OPTARG" ;;
         o ) output_dir="$OPTARG" ;;
         t ) timestamp="$OPTARG" ;;
-        g ) gpu_list="$OPTARG" ;;
+        s ) single=true ;;
         ? ) helpFunction_launch_eval ;; 
     esac
 done
@@ -37,7 +41,7 @@ shift $((OPTIND-1))
 
 # Deal with gpu's. If passed in, use those.
 GPU_IDX=("$@")
-if [ -z "$GPU_IDX" ]; then
+if [ -z "${GPU_IDX[0]+x}" ]; then
     echo "no gpus set... finding available gpus"
     # Find available devices
     num_device=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
@@ -45,32 +49,34 @@ if [ -z "$GPU_IDX" ]; then
     END=${num_device}-1
     GPU_IDX=()
 
-    for (( id=$START; id<$END; id++ )); do
-        free_mem=$(nvidia-smi --query-gpu=memory.free --format=csv -i $id | grep -Eo [0-9]+)
+    for (( id=START; id<=END; id++ )); do
+        free_mem=$(nvidia-smi --query-gpu=memory.free --format=csv -i $id | grep -Eo '[0-9]+')
         if [[ $free_mem -gt 10000 ]]; then
-            GPU_IDX+=( $id )
+            GPU_IDX+=( "$id" )
         fi
     done
 fi
-echo available gpus... ${GPU_IDX[@]}
+echo "available gpus... ${GPU_IDX[*]}"
 
 DATASETS=("mic" "ficus" "chair" "hotdog" "materials" "drums" "ship" "lego")
 idx=0
 len=${#GPU_IDX[@]}
-((len=len-1))
+GPU_PID=()
+# kill all the background jobs if terminated:
+trap "trap - SIGTERM && kill -- -$$" SIGINT SIGTERM EXIT
 
-for dataset in ${DATASETS[@]}; do
+for dataset in "${DATASETS[@]}"; do
+    if "$single" && [ -n "${GPU_PID[$idx]+x}" ]; then
+        wait "${GPU_PID[$idx]}"
+    fi
     export CUDA_VISIBLE_DEVICES=${GPU_IDX[$idx]}
     config_path="${output_dir}/blender_${dataset}_${timestamp::-7}/${method_name}/${timestamp}/config.yml"
-    ns-eval --load-config=${config_path} \
-            --output-path=${output_dir}/${method_name}/blender_${dataset}_${timestamp}.json &
+    ns-eval --load-config="${config_path}" \
+            --output-path="${output_dir}/${method_name}/blender_${dataset}_${timestamp}.json" & GPU_PID[$idx]=$!
     echo "Launched ${config_path} on gpu ${GPU_IDX[$idx]}"
 
     # update gpu
-    if [ $idx == $len ]; then
-        idx=0
-    else
-        ((idx=idx+1))
-    fi
-
+    ((idx=(idx+1)%len))
 done
+wait
+echo "Done."

scripts/benchmarking/launch_train_blender.sh

@@ -2,30 +2,40 @@
 
 helpFunction_launch_train()
 {
-   echo "Usage: $0 -m method_name -g [OPTIONAL] gpu_list"
+   echo "Usage: $0 -m <method_name> [-v <vis>] [-s] [<gpu_list>]"
    echo -e "\t-m name of config to benchmark (e.g. mipnerf, instant_ngp)"
-   echo -e "\t-g [OPTIONAL] list of space-separated gpu numbers to launch train on (e.g. 0 2 4 5)"
+   echo -e "\t-v <vis>: Visualization method. <vis> can be wandb or tensorboard. Default is wandb."
+   echo -e "\t-s: Launch a single training job per gpu."
+   echo -e "\t<gpu_list> [OPTIONAL] list of space-separated gpu numbers to launch train on (e.g. 0 2 4 5)"
    exit 1 # Exit program after printing help
 }
 
-while getopts "m:g" opt; do
+vis="wandb"
+single=false
+while getopts "m:v:s" opt; do
     case "$opt" in
         m ) method_name="$OPTARG" ;;
-        g ) gpu_list="$OPTARG" ;;
+        v ) vis="$OPTARG" ;;
+        s ) single=true ;;
         ? ) helpFunction ;; 
     esac
 done
 
-if [ -z "$method_name" ]; then
+if [ -z "${method_name+x}" ]; then
     echo "Missing method name"
     helpFunction_launch_train
 fi
+method_opts=()
+if [ "$method_name" = "nerfacto" ]; then
+    # https://github.com/nerfstudio-project/nerfstudio/issues/806#issuecomment-1284327844
+    method_opts=(--pipeline.model.near-plane 2. --pipeline.model.far-plane 6. --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-average-appearance-embedding False)
+fi
 
 shift $((OPTIND-1))
 
 # Deal with gpu's. If passed in, use those.
 GPU_IDX=("$@")
-if [ -z "$GPU_IDX" ]; then
+if [ -z "${GPU_IDX[0]+x}" ]; then
     echo "no gpus set... finding available gpus"
     # Find available devices
     num_device=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
@@ -33,41 +43,51 @@ if [ -z "$GPU_IDX" ]; then
     END=${num_device}-1
     GPU_IDX=()
 
-    for (( id=$START; id<$END; id++ )); do
-        free_mem=$(nvidia-smi --query-gpu=memory.free --format=csv -i $id | grep -Eo [0-9]+)
+    for (( id=START; id<=END; id++ )); do
+        free_mem=$(nvidia-smi --query-gpu=memory.free --format=csv -i $id | grep -Eo '[0-9]+')
         if [[ $free_mem -gt 10000 ]]; then
-            GPU_IDX+=( $id )
+            GPU_IDX+=( "$id" )
         fi
     done
 fi
-echo available gpus... ${GPU_IDX[@]}
+echo "available gpus... ${GPU_IDX[*]}"
 
 DATASETS=("mic" "ficus" "chair" "hotdog" "materials" "drums" "ship" "lego")
 date
-now=$(date)
 tag=$(date +'%Y-%m-%d')
 idx=0
 len=${#GPU_IDX[@]}
-((len=len-1))
+GPU_PID=()
+timestamp=$(date "+%Y-%m-%d_%H%M%S")
+# kill all the background jobs if terminated:
+trap "trap - SIGTERM && kill -- -$$" SIGINT SIGTERM EXIT
 
-for dataset in ${DATASETS[@]}; do
-    export CUDA_VISIBLE_DEVICES=${GPU_IDX[$idx]}
-    ns-train ${method_name} \
-             --data=data/blender/${dataset} \
-             --experiment-name=blender_${dataset}_${tag} \
+for dataset in "${DATASETS[@]}"; do
+    if "$single" && [ -n "${GPU_PID[$idx]+x}" ]; then
+        echo "Waiting for GPU ${GPU_IDX[$idx]}"
+        wait "${GPU_PID[$idx]}"
+        echo "GPU ${GPU_IDX[$idx]} is available"
+    fi
+    export CUDA_VISIBLE_DEVICES="${GPU_IDX[$idx]}"
+    ns-train "${method_name}" "${method_opts[@]}" \
+             --data="data/blender/${dataset}" \
+             --experiment-name="blender_${dataset}_${tag}" \
              --trainer.relative-model-dir=nerfstudio_models/ \
              --trainer.steps-per-save=1000 \
              --trainer.max-num-iterations=16500 \
-             --logging.local-writer.no-enable  \
-             --logging.no-enable-profiler \
-             --vis wandb \
-             blender-data &
+             --logging.local-writer.enable=False  \
+             --logging.enable-profiler=False \
+             --vis "${vis}" \
+             --timestamp "$timestamp" \
+             blender-data & GPU_PID[$idx]=$!
     echo "Launched ${method_name} ${dataset} on gpu ${GPU_IDX[$idx]}, ${tag}"
     
     # update gpu
-    if [ $idx == $len ]; then
-        idx=0
-    else
-        ((idx=idx+1))
-    fi
-done
+    ((idx=(idx+1)%len))
+done
+wait
+echo "Done."
+echo "Launch eval with:"
+s=""
+$single && s="-s"
+echo "$(dirname "$0")/launch_eval_blender.sh -m $method_name -o outputs/ -t $timestamp $s"