Sync run_eval with benchmarking script (#260)

* added new .sh script for eval; removed old python script

* fixed changing dataset

* done refactoring eval script

scripts/benchmarking/launch_eval_blender.sh

@@ -0,0 +1,76 @@
+helpFunction_launch_eval()
+{
+   echo "Usage: $0 -c config_name -g gpu_list"
+   echo -e "\t-c name of method to benchmark (e.g. graph_mipnerf, graph_instant_ngp)"
+   echo -o "\t-o base directory for where all the benchmarks are stored (e.g. outputs/)"
+   echo -e "\t-m month of the benchmark of format xx (e.g. 08)"
+   echo -e "\t-d date of the benchmark of format xx (e.g. 01)"
+   echo -e "\t-y year of the benchmark of format xxxx (e.g. 2022)"
+   echo -e "\t-s second timestamp of the benchmark of format xxxxxx (e.g. 003603)"
+   echo -e "\t-g 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 "c:o:m:d:y:s:g" opt; do
+    case "$opt" in
+        c ) config_name="$OPTARG" ;;
+        o ) output_dir="$OPTARG" ;;
+        m ) month="$OPTARG" ;;
+        d ) date="$OPTARG" ;;
+        y ) year="$OPTARG" ;;
+        s ) seconds="$OPTARG" ;;
+        g ) gpu_list="$OPTARG" ;;
+        ? ) helpFunction_launch_eval ;; 
+    esac
+done
+
+if [ -z "$config_name" ]; then
+    echo "Missing method name"
+    helpFunction_launch_train
+fi
+
+shift $((OPTIND-1))
+
+# Deal with gpu's. If passed in, use those.
+GPU_IDX=("$@")
+if [ -z "$GPU_IDX" ]; then
+    echo "no gpus set... finding available gpus"
+    # Find available devices
+    num_device=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
+    START=0
+    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]+)
+        if [[ $free_mem -gt 10000 ]]; then
+            GPU_IDX+=( $id )
+        fi
+    done
+fi
+echo available gpus... ${GPU_IDX[@]}
+
+DATASETS=("mic" "ficus" "chair" "hotdog" "materials" "drums" "ship" "lego")
+idx=0
+len=${#GPU_IDX[@]}
+((len=len-1))
+
+for dataset in ${DATASETS[@]}; do
+    export CUDA_VISIBLE_DEVICES=${GPU_IDX[$idx]}
+    base_config_name=${config_name/"graph_"/""}
+    checkpoint_dir="${output_dir}/blender_${dataset}_${month}-${date}-${year}/${base_config_name}/${year}-${month}-${date}_${seconds}/nerfactory_models/"
+    python scripts/run_eval.py --config-name ${config_name} \
+        --checkpoint-dir ${checkpoint_dir} \
+        --rendered-output-name rgb \
+        --output-filename ${output_dir}/${base_config_name}/blender_${dataset}_${month}-${date}-${year}_${seconds}.json \
+        --config-overrides data.dataset_inputs_train.data_directory=data/blender/${dataset} data.dataset_inputs_eval.data_directory=data/blender/${dataset} &
+    echo "Launched ${checkpoint_dir} on gpu ${GPU_IDX[$idx]}"
+
+    # update gpu
+    if [ $idx == $len ]; then
+        idx=0
+    else
+        ((idx=idx+1))
+    fi
+
+done

scripts/benchmarking/launch_train_blender.sh

@@ -1,7 +1,27 @@
 #!/bin/bash
 
-METHOD=$1
-shift
+helpFunction_launch_train()
+{
+   echo "Usage: $0 -c config_name -g gpu_list"
+   echo -e "\t-c name of config to benchmark (e.g. graph_mipnerf, graph_instant_ngp)"
+   echo -e "\t-g 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 "c:g" opt; do
+    case "$opt" in
+        c ) config_name="$OPTARG" ;;
+        g ) gpu_list="$OPTARG" ;;
+        ? ) helpFunction ;; 
+    esac
+done
+
+if [ -z "$config_name" ]; then
+    echo "Missing method name"
+    helpFunction_launch_train
+fi
+
+shift $((OPTIND-1))
 
 # Deal with gpu's. If passed in, use those.
 GPU_IDX=("$@")
@@ -32,13 +52,8 @@ len=${#GPU_IDX[@]}
 
 for dataset in ${DATASETS[@]}; do
     export CUDA_VISIBLE_DEVICES=${GPU_IDX[$idx]}
-    if [ $idx == $len ]; then
-        idx=0
-    else
-        ((idx=idx+1))
-    fi
     python scripts/run_train.py \
-           --config-name ${METHOD} \
+           --config-name ${config_name} \
            '~logging.writer.LocalWriter' \
            data.dataset_inputs_train.data_directory=data/blender/${dataset} \
            data.dataset_inputs_eval.data_directory=data/blender/${dataset} \
@@ -48,5 +63,12 @@ for dataset in ${DATASETS[@]}; do
            trainer.max_num_iterations=2000000 \
            viewer.enable=False \
            logging.enable_profiler=False &
-    echo "Launched ${METHOD} ${dataset} on gpu ${GPU_IDX[$idx]}, ${tag}"
+    echo "Launched ${config_name} ${dataset} on gpu ${GPU_IDX[$idx]}, ${tag}"
+    
+    # update gpu
+    if [ $idx == $len ]; then
+        idx=0
+    else
+        ((idx=idx+1))
+    fi
 done

scripts/benchmarking/run_benchmark.py

@@ -1,150 +0,0 @@
-"""
-run_benchmark.py
-"""
-import argparse
-import json
-import logging
-import os
-from datetime import datetime
-from typing import Any, Dict, Tuple
-
-import numpy as np
-from hydra import compose, initialize
-from omegaconf import DictConfig
-
-from scripts.run_eval import run_inference_from_config
-
-
-def _load_best_ckpt(hydra_dir: str, config: DictConfig) -> str:
-    """helper function to update config with latest checkpoint in specified hydra_dir
-
-    Args:
-        hydra_dir (str): base directory for the specified run (where the model directories are stored)
-        config (DictConfig): configuration for the specified run.
-
-    Returns:
-        str: path to the most recent checkpoint sorted by timestamps and step number
-    """
-    # get the most recent run in the model directory
-    latest_run = os.listdir(hydra_dir)
-    latest_run.sort()
-    latest_run = latest_run[-1]
-    # get the latest checkpoint name
-    model_dir = os.path.join(hydra_dir, latest_run, config.model_dir)
-    latest_ckpt = os.listdir(model_dir)
-    latest_ckpt.sort()
-    latest_ckpt = latest_ckpt[-1]
-    step = os.path.splitext(latest_ckpt)[0].split("-")[-1]
-    config.resume_train.load_dir = model_dir
-    config.resume_train.load_step = int(step)
-    return os.path.join(model_dir, latest_ckpt)
-
-
-def _load_hydra_config(hydra_dir: str, benchmark_time=None, benchmark_date=None) -> Tuple[DictConfig, str]:
-    """helper function to load the specified hydra config from the specified directory
-
-    Args:
-        hydra_dir (str): base directory for the specified run (where the flattened config.yaml is stored)
-
-    Returns:
-        DictConfig: returns the loaded hydra dictionary config
-    """
-    if benchmark_time:
-        assert benchmark_date is not None
-        reformat_date = datetime.strptime(benchmark_date, "%m-%d-%Y").strftime("%Y-%m-%d")
-        basename = f"{reformat_date}_{benchmark_time}"
-    else:
-        basename = sorted(os.listdir(hydra_dir))[-1]
-    hydra_dir = f"{hydra_dir}/{basename}"
-    initialize(version_base="1.2", config_path=os.path.join("../../", hydra_dir, ".hydra/"))
-    config = compose("config.yaml")
-    return config, basename
-
-
-def _calc_avg(stat_name: str, benchmark: Dict[str, Any]):
-    """helper to calculate the average across all objects in dataset"""
-    stats = []
-    for _, object_stats in benchmark.items():
-        stats.append(object_stats[stat_name])
-    return np.mean(stats)
-
-
-def main(args):
-    """Main function."""
-    benchmarks = {}
-    # set up trainer, config, and checkpoint loading
-    hydra_dir = f"{args.hydra_base_dir}/blender_{args.item_name}_{args.benchmark_date}/{args.graph}/"
-    config, basename = _load_hydra_config(
-        hydra_dir, benchmark_time=args.benchmark_time, benchmark_date=args.benchmark_date
-    )
-    config.data.dataloader_eval.image_indices = None
-    ckpt = _load_best_ckpt(hydra_dir, config.trainer)
-
-    # run evaluation
-    stats_dict = run_inference_from_config(config)
-    stats_dict["checkpoint"] = ckpt
-    benchmarks[args.item_name] = stats_dict
-
-    avg_rays_per_sec = _calc_avg("avg rays per sec", benchmarks)
-    avg_fps = _calc_avg("avg fps", benchmarks)
-
-    # output benchmark statistics to a json file
-    benchmark_info = {
-        "meta info": {"graph": args.graph, "benchmark_date": args.benchmark_date, "hydra": args.hydra_base_dir},
-        "avg rays per sec": avg_rays_per_sec,
-        "avg fps": avg_fps,
-        "results": benchmarks,
-    }
-    json_dir = os.path.join(args.hydra_base_dir, args.graph)
-    if not os.path.exists(json_dir):
-        os.makedirs(json_dir)
-    json_file = os.path.join(json_dir, f"{basename}_{args.item_name}.json")
-
-    # if benchmark file already exists, update the results dictionary by overriding existing info or appending.
-    if os.path.exists(json_file):
-        with open(json_file, "r", encoding="utf8") as f:
-            existing_data = json.load(f)
-        existing_benchmarks = existing_data["results"]
-        existing_benchmarks.update(benchmarks)
-        existing_data["results"] = existing_benchmarks
-        benchmark_info = existing_data
-
-    with open(json_file, "w", encoding="utf8") as f:
-        json.dump(benchmark_info, f, indent=2)
-    logging.info("saved benchmark results to %s", json_file)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "-g",
-        "--graph",
-        type=str,
-        required=True,
-        help="name of nerf graph to benchmark. Note: name without the 'graph_' prefix",
-    )
-    parser.add_argument(
-        "-d",
-        "--benchmark_date",
-        type=str,
-        required=True,
-        help="date of run to benchmark provided in '%%m-%%d-%%Y' format.",
-    )
-    parser.add_argument(
-        "-i",
-        "--item_name",
-        type=str,
-        required=True,
-        choices=["mic", "ficus", "chair", "hotdog", "materials", "drums", "ship", "lego"],
-        help="name of item in blender dataset to benchmark",
-    )
-    parser.add_argument(
-        "-t",
-        "--benchmark_time",
-        type=str,
-        default=None,
-        help="timestamp of the run to benchmark; if None, will default to running most recent timestamp",
-    )
-    parser.add_argument("-o", "--hydra_base_dir", type=str, default="outputs/", help="hydra base output path")
-    parsed_args = parser.parse_args()
-    main(parsed_args)

scripts/debugging/profile.sh

@@ -27,6 +27,7 @@ done
 
 if [ -z "$type" ]; then
     echo "Missing profiling type (flame / top)"
+    exit 1
 fi
 shift $((OPTIND-1))
 program="$@"

scripts/run_eval.py

@@ -2,6 +2,7 @@
 run_eval.py
 """
 import enum
+import json
 import os
 from typing import Any, Dict, List, Optional
 
@@ -41,7 +42,7 @@ def _update_avg(prev_avg: float, new_val: float, step: int) -> float:
     return (step * prev_avg + new_val) / (step + 1)
 
 
-def _load_checkpoint(config: DictConfig, graph: Graph) -> None:
+def _load_checkpoint(config: DictConfig, graph: Graph) -> str:
     """Helper function to load checkpointed graph
 
     Args:
@@ -60,6 +61,7 @@ def _load_checkpoint(config: DictConfig, graph: Graph) -> None:
     loaded_state = torch.load(load_path, map_location="cpu")
     graph.load_graph(loaded_state)
     print(f"done loading checkpoint from {load_path}")
+    return load_path
 
 
 def run_inference_from_config(config: DictConfig) -> Dict[str, Any]:
@@ -80,7 +82,7 @@ def run_inference_from_config(config: DictConfig) -> Dict[str, Any]:
     graph.eval()
 
     # load checkpointed information
-    _load_checkpoint(config.trainer.resume_train, graph)
+    _ = _load_checkpoint(config.trainer.resume_train, graph)
 
     # calculate average psnr across test dataset
     return render_stats_dict(graph, dataloader_eval)
@@ -202,7 +204,7 @@ def main(
     graph = setup_graph(config.graph, dataset_inputs_train, device=device)
     graph.eval()
     # load checkpointed information
-    _load_checkpoint(config.trainer.resume_train, graph)
+    checkpoint_name = _load_checkpoint(config.trainer.resume_train, graph)
 
     if method == MethodType.PSNR:
         stats_dict = render_stats_dict(graph, dataloader_eval)
@@ -212,6 +214,19 @@ def main(
         print(f"Avg. PSNR: {avg_psnr:0.4f}")
         print(f"Avg. Rays / Sec: {human_format(avg_rays_per_sec)}")
         print(f"Avg. FPS: {avg_fps:0.4f}")
+        # save output to some file
+        dirname = os.path.dirname(output_filename)
+        if not os.path.exists(dirname):
+            os.makedirs(dirname)
+        benchmark_info = {
+            "config_name": config_name,
+            "checkpoint": checkpoint_name,
+            "results": stats_dict,
+        }
+        with open(output_filename, "w", encoding="utf8") as f:
+            json.dump(benchmark_info, f, indent=2)
+        print(f"Saved results to: {output_filename}")
+
     elif method == MethodType.TRAJ:
         # TODO(ethan): pass in camera information into argparse parser
         if traj == TrajectoryType.SPIRAL: