From 6a70eaa40a64d574fc0a7ba90e1101afef47553b Mon Sep 17 00:00:00 2001
From: evonneng <evonne_ng@berkeley.edu>
Date: Thu, 12 May 2022 17:32:53 +0000
Subject: [PATCH] fixed up benchmarking test script more versatile

---
 README.md                                    |  6 +++
 scripts/benchmarking/launch_train_blender.sh |  1 +
 scripts/benchmarking/run_benchmark.py        | 45 +++++++++++++-------
 3 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/README.md b/README.md
index 03f34916c..c96306876 100644
--- a/README.md
+++ b/README.md
@@ -110,4 +110,10 @@ ipython --pdb scripts/run_train.py
 
 ```
 ./scripts/benchmarking/launch_train_blender.sh <gpu0> <gpu1> ... <gpu7>
+```
+
+For testing specific methods, see scripts/benchmarking/run_benchmark.py.
+Modify the `BENCH` variable to specify which jobs ("ckpt_dir") and methods ("method") you want benchmarked. Then run
+```
+python scripts/benchmarking/run_benchmark.py
 ```
\ No newline at end of file
diff --git a/scripts/benchmarking/launch_train_blender.sh b/scripts/benchmarking/launch_train_blender.sh
index 803a0be76..125d3c7cd 100755
--- a/scripts/benchmarking/launch_train_blender.sh
+++ b/scripts/benchmarking/launch_train_blender.sh
@@ -6,6 +6,7 @@ for dataset in ${DATASETS[@]}; do
     export CUDA_VISIBLE_DEVICES=${!GPU_IDX}
     python scripts/run_train.py \
            data.dataset.data_directory=data/blender/${dataset} \
+           data.dataset.downscale_factor=1 \
            experiment_name=blender_${dataset} \
            graph.model_dir=mattport_models/ \
            graph.steps_per_save=25000 \
diff --git a/scripts/benchmarking/run_benchmark.py b/scripts/benchmarking/run_benchmark.py
index 9c2a5d43d..d41dbc674 100644
--- a/scripts/benchmarking/run_benchmark.py
+++ b/scripts/benchmarking/run_benchmark.py
@@ -2,27 +2,35 @@
 run_benchmark.py
 """
 import json
+import logging
 import os
+from datetime import date
 
 import torch
 from hydra import compose, initialize
+from hydra.core.global_hydra import GlobalHydra
 from omegaconf import DictConfig
+from tqdm import tqdm
 
 from mattport.nerf.trainer import Trainer
 
 BENCH = {
-    "mic": "outputs/blender_mic/vanilla_nerf/2022-05-11_192109/",
-    "ficus": "outputs/blender_ficus/vanilla_nerf/2022-05-11_192109/",
-    "chair": "outputs/blender_chair/vanilla_nerf/2022-05-11_192109/",
-    "hotdog": "outputs/blender_hotdog/vanilla_nerf/2022-05-11_192109/",
-    "materials": "outputs/blender_materials/vanilla_nerf/2022-05-11_192109/",
-    "drums": "outputs/blender_drums/vanilla_nerf/2022-05-11_192109/",
-    "ship": "outputs/blender_ship/vanilla_nerf/2022-05-11_192109/",
-    "lego": "outputs/blender_lego/vanilla_nerf/2022-05-11_192109/",
+    "method": "vanilla_nerf",
+    "hydra_base_dir": "outputs/",
+    "ckpt_dir": {
+        "mic": "2022-05-11_192109",
+        "ficus": "2022-05-11_192109",
+        "chair": "2022-05-11_192109",
+        "hotdog": "2022-05-11_192109",
+        "materials": "2022-05-11_192109",
+        "drums": "2022-05-11_192109",
+        "ship": "2022-05-11_192109",
+        "lego": "2022-05-11_192109",
+    },
 }
 
 
-def load_best_ckpt(hydra_dir: str, config: DictConfig):
+def load_best_ckpt(hydra_dir: str, config: DictConfig) -> str:
     """helper function to update config with latest checkpoint in specified hydra_dir"""
     model_dir = os.path.join(hydra_dir, config.graph.model_dir)
     latest_ckpt = os.listdir(model_dir)
@@ -32,15 +40,19 @@ def load_best_ckpt(hydra_dir: str, config: DictConfig):
     config.graph.resume_train.load_dir = model_dir
     config.graph.resume_train.load_step = int(step)
     config.logging.writer.TensorboardWriter.save_dir = "/tmp/"
+    return os.path.join(model_dir, latest_ckpt)
 
 
 def main():
     """Main function."""
     benchmarks = {}
-    for dataset, hydra_dir in BENCH.items():
+    hydra_base_dir = BENCH["hydra_base_dir"]
+    method = BENCH["method"]
+    for dataset, ckpt_dir in tqdm(BENCH["ckpt_dir"].items()):
+        hydra_dir = f"{hydra_base_dir}/blender_{dataset}/{method}/{ckpt_dir}/"
         initialize(config_path=os.path.join("../../", hydra_dir, ".hydra/"))
         config = compose("config.yaml")
-        load_best_ckpt(hydra_dir, config)
+        ckpt = load_best_ckpt(hydra_dir, config)
         trainer = Trainer(config, local_rank=0, world_size=1)
         trainer.setup(test_mode=True)
         avg_psnr = 0
@@ -48,12 +60,15 @@ def main():
             with torch.no_grad():
                 psnr = trainer.test_image(image_idx=image_idx, step=step)
             avg_psnr = (step * avg_psnr + psnr) / (step + 1)
-        benchmarks[dataset] = avg_psnr
+        benchmarks[dataset] = (avg_psnr, ckpt)
+        GlobalHydra.instance().clear()
 
     benchmark_info = {"bench": BENCH, "results": benchmarks}
-    # TODO(more sophisticated saving to output directory)
-    with open("/tmp/benchmark.json", "w", encoding="utf8") as f:
-        json.dump(benchmark_info, f)
+    timestamp = date.today().strftime("%b-%d-%Y")
+    json_file = os.path.join(BENCH["hydra_base_dir"], f"{timestamp}.json")
+    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__":