From 1dc574281ca2c605c5e2ce0c270195ad53a6ea40 Mon Sep 17 00:00:00 2001 From: evonneng Date: Wed, 11 May 2022 17:56:39 -0700 Subject: [PATCH] Added eval script for benchmarking blender dataset (#71) * fixed train iteration to much longer * done adding test benchmark script --- mattport/nerf/graph/vanilla_nerf.py | 1 + mattport/nerf/trainer.py | 71 +++++++++++--------- scripts/benchmarking/launch_train_blender.sh | 3 + scripts/benchmarking/run_benchmark.py | 60 +++++++++++++++++ 4 files changed, 104 insertions(+), 31 deletions(-) create mode 100644 scripts/benchmarking/run_benchmark.py diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index ce8c2bcf3..b35d737f4 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -246,3 +246,4 @@ class NeRFGraph(Graph): writer.write_event( {"name": f"val_idx_{image_idx}", "scalar": float(fine_lpips), "step": step, "group": "lpips"} ) + return fine_psnr.item() diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index e1d149d3e..cf2fa468c 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -56,10 +56,10 @@ class Trainer: self.device = "cpu" if self.world_size == 0 else f"cuda:{self.local_rank}" @profiler.time_function - def setup(self): + def setup(self, test_mode=False): """Setup the Trainer by calling other setup functions.""" dataset_inputs_dict = get_dataset_inputs_dict(**self.config.data.dataset) - self.setup_datasets(dataset_inputs_dict) + self.setup_datasets(dataset_inputs_dict, test_mode=test_mode) self.setup_graph(dataset_inputs_dict["train"]) def collate_fn(self, batch_list): @@ -68,35 +68,43 @@ class Trainer: """ return collate_batch(batch_list, self.config.data.dataloader.num_rays_per_batch, keep_full_image=False) - @profiler.time_function - def setup_datasets(self, dataset_inputs_dict: Dict[str, DatasetInputs]): - """_summary_""" - self.train_image_dataset = ImageDataset( - image_filenames=dataset_inputs_dict["train"].image_filenames, - downscale_factor=dataset_inputs_dict["train"].downscale_factor, - alpha_color=dataset_inputs_dict["train"].alpha_color, - ) - - self.train_dataset = CollateIterDataset( - self.train_image_dataset, - collate_fn=self.collate_fn, - num_samples_to_collate=self.config.data.dataloader.num_images_to_sample_from, - num_times_to_repeat=self.config.data.dataloader.num_times_to_repeat_images, - ) - self.train_dataloader = DataLoader( - self.train_dataset, - batch_size=1, - num_workers=self.config.data.dataloader.num_workers, - collate_fn=collate_batch_size_one, - pin_memory=True, - ) + def load_test_dataset(self, dataset_inputs_dict, mode): + """Helper method to load test or val dataset based on test/train mode""" self.val_image_dataset = ImageDataset( - image_filenames=dataset_inputs_dict["val"].image_filenames, - downscale_factor=dataset_inputs_dict["val"].downscale_factor, - alpha_color=dataset_inputs_dict["val"].alpha_color, + image_filenames=dataset_inputs_dict[mode].image_filenames, + downscale_factor=dataset_inputs_dict[mode].downscale_factor, + alpha_color=dataset_inputs_dict[mode].alpha_color, ) - self.val_image_intrinsics = dataset_inputs_dict["val"].intrinsics - self.val_image_camera_to_world = dataset_inputs_dict["val"].camera_to_world + self.val_image_intrinsics = dataset_inputs_dict[mode].intrinsics + self.val_image_camera_to_world = dataset_inputs_dict[mode].camera_to_world + + @profiler.time_function + def setup_datasets(self, dataset_inputs_dict: Dict[str, DatasetInputs], test_mode: bool): + """_summary_""" + if test_mode: + # TODO(change "val" to "test") + self.load_test_dataset(dataset_inputs_dict, "val") + else: + self.train_image_dataset = ImageDataset( + image_filenames=dataset_inputs_dict["train"].image_filenames, + downscale_factor=dataset_inputs_dict["train"].downscale_factor, + alpha_color=dataset_inputs_dict["train"].alpha_color, + ) + + self.train_dataset = CollateIterDataset( + self.train_image_dataset, + collate_fn=self.collate_fn, + num_samples_to_collate=self.config.data.dataloader.num_images_to_sample_from, + num_times_to_repeat=self.config.data.dataloader.num_times_to_repeat_images, + ) + self.train_dataloader = DataLoader( + self.train_dataset, + batch_size=1, + num_workers=self.config.data.dataloader.num_workers, + collate_fn=collate_batch_size_one, + pin_memory=True, + ) + self.load_test_dataset(dataset_inputs_dict, "val") @profiler.time_function def setup_graph(self, dataset_inputs: DatasetInputs): @@ -216,7 +224,7 @@ class Trainer: self.save_checkpoint(self.config.graph.model_dir, step) if step % self.config.graph.steps_per_test == 0: for image_idx in self.config.data.validation_image_indices: - self.test_image(image_idx=image_idx, step=step) + _ = self.test_image(image_idx=image_idx, step=step) stats_tracker.print_stats(i / num_iterations) stats_tracker.update_stats( @@ -245,5 +253,6 @@ class Trainer: camera_to_world = self.val_image_camera_to_world[image_idx] outputs = self.graph.get_outputs_for_camera(intrinsics, camera_to_world) image = self.val_image_dataset[image_idx]["image"].to(self.device) - self.graph.log_test_image_outputs(image_idx, step, image, outputs) + psnr = self.graph.log_test_image_outputs(image_idx, step, image, outputs) self.graph.train() + return psnr diff --git a/scripts/benchmarking/launch_train_blender.sh b/scripts/benchmarking/launch_train_blender.sh index 31165ea4c..803a0be76 100755 --- a/scripts/benchmarking/launch_train_blender.sh +++ b/scripts/benchmarking/launch_train_blender.sh @@ -7,6 +7,9 @@ for dataset in ${DATASETS[@]}; do python scripts/run_train.py \ data.dataset.data_directory=data/blender/${dataset} \ experiment_name=blender_${dataset} \ + graph.model_dir=mattport_models/ \ + graph.steps_per_save=25000 \ + graph.max_num_iterations=2000000 \ logging.enable_stats=False \ logging.enable_profiler=False & echo "Launched ${dataset} on gpu ${!GPU_IDX}" diff --git a/scripts/benchmarking/run_benchmark.py b/scripts/benchmarking/run_benchmark.py new file mode 100644 index 000000000..9c2a5d43d --- /dev/null +++ b/scripts/benchmarking/run_benchmark.py @@ -0,0 +1,60 @@ +""" +run_benchmark.py +""" +import json +import os + +import torch +from hydra import compose, initialize +from omegaconf import DictConfig + +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/", +} + + +def load_best_ckpt(hydra_dir: str, config: DictConfig): + """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) + latest_ckpt.sort() + latest_ckpt = latest_ckpt[-1] + step = os.path.splitext(latest_ckpt)[0].split("-")[-1] + config.graph.resume_train.load_dir = model_dir + config.graph.resume_train.load_step = int(step) + config.logging.writer.TensorboardWriter.save_dir = "/tmp/" + + +def main(): + """Main function.""" + benchmarks = {} + for dataset, hydra_dir in BENCH.items(): + initialize(config_path=os.path.join("../../", hydra_dir, ".hydra/")) + config = compose("config.yaml") + load_best_ckpt(hydra_dir, config) + trainer = Trainer(config, local_rank=0, world_size=1) + trainer.setup(test_mode=True) + avg_psnr = 0 + for step, image_idx in enumerate(config.data.validation_image_indices): + 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 + + 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) + + +if __name__ == "__main__": + main()