tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

fixed up benchmarking test script more versatile

Browse Source
This commit is contained in:
evonneng 2022-05-12 17:32:53 +00:00
parent 39db2e8a27
commit 6a70eaa40a
3 changed files with 37 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
README.md
View File

@ -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
```

1
scripts/benchmarking/launch_train_blender.sh
View File

@ -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 \

45
scripts/benchmarking/run_benchmark.py
View File

@ -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__":