From 023e6a3f92085df53850ccba56992d5cfab51a43 Mon Sep 17 00:00:00 2001
From: Ethan Weber <eweb0124@gmail.com>
Date: Fri, 29 Apr 2022 22:29:43 +0000
Subject: [PATCH] merging masters new config

---
 configs/data/default_data.yaml |  3 +++
 mattport/nerf/trainer.py       | 19 ++++++++++---------
 scripts/run_train.py           |  2 +-
 tests/configs/test_default.yml |  3 +++
 4 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/configs/data/default_data.yaml b/configs/data/default_data.yaml
index 736eea12e..a42faf749 100644
--- a/configs/data/default_data.yaml
+++ b/configs/data/default_data.yaml
@@ -4,6 +4,9 @@ dataset:
   dataset_type: blender
   downscale_factor: 4
 
+# validation options
+validation_image_indices: [0, 10] # Which images to evaluate on.
+
 # dataloader options
 dataloader:
   num_workers: 4
diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py
index 9f547e05f..490e48b9b 100644
--- a/mattport/nerf/trainer.py
+++ b/mattport/nerf/trainer.py
@@ -52,12 +52,11 @@ class Trainer:
         self.start_step = 0
         # logging variables
         self.is_main_thread = world_size != 0 and local_rank % world_size == 0
-        writer = getattr(mattport.utils.writer, self.config.logging_configs.writer.type)
-        self.writer = writer(self.is_main_thread, save_dir=self.config.logging_configs.writer.save_dir)
+        writer = getattr(mattport.utils.writer, self.config.logging.writer.type)
+        self.writer = writer(self.is_main_thread, save_dir=self.config.logging.writer.save_dir)
         self.stats = StatsTracker(config, self.is_main_thread)
         if not profiler.PROFILER and self.config.logging.enable_profiler:
             profiler.PROFILER = profiler.Profiler(config, self.is_main_thread)
-        self.dry_run = self.config.get("dry_run", False)
         self.device = f"cuda:{self.local_rank}" if not cpu else "cpu"
 
     @profiler.time_function
@@ -71,7 +70,7 @@ class Trainer:
         """TODO(ethan): I need to replace this.
         I'm only using this for multiprocess pickle issues for now.
         """
-        return collate_batch(batch_list, self.config.dataloader.num_rays_per_batch, keep_full_image=False)
+        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]):
@@ -84,8 +83,8 @@ class Trainer:
         self.train_dataset = CollateIterDataset(
             self.train_image_dataset,
             collate_fn=self.collate_fn,
-            num_samples_to_collate=self.config.dataloader.num_images_to_sample_from,
-            num_times_to_repeat=self.config.dataloader.num_times_to_repeat_images,
+            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,
@@ -107,7 +106,9 @@ class Trainer:
             dataset_inputs (DatasetInputs): The inputs which will be used to define the camera parameters.
         """
         self.graph = instantiate(
-            self.config.graph, intrinsics=dataset_inputs.intrinsics, camera_to_world=dataset_inputs.camera_to_world
+            self.config.graph.network,
+            intrinsics=dataset_inputs.intrinsics,
+            camera_to_world=dataset_inputs.camera_to_world,
         ).to(self.device)
         self.setup_optimizers()
 
@@ -199,8 +200,8 @@ class Trainer:
             if step != 0 and self.config.graph.steps_per_save and step % self.config.graph.steps_per_save == 0:
                 self.save_checkpoint(self.config.graph.model_dir, step)
             if step % self.config.graph.steps_per_test == 0:
-                self.test_image(image_idx=0, step=step)
-                _ = self.test_image(image_idx=10, step=step) if not self.dry_run else None
+                for image_idx in self.config.data.validation_image_indices:
+                    self.test_image(image_idx=image_idx, step=step)
             self.stats.print_stats(i / num_iterations)
 
         self.stats.update_time(Stats.TOTAL_TRAIN_TIME, train_start, time(), step=-1)
diff --git a/scripts/run_train.py b/scripts/run_train.py
index 3a7c77515..28c587a15 100644
--- a/scripts/run_train.py
+++ b/scripts/run_train.py
@@ -153,7 +153,7 @@ def launch(
         except KeyboardInterrupt:
             pass
         finally:
-            if config.logging_configs.enable_profiler:
+            if config.logging.enable_profiler:
                 profiler.PROFILER.print_profile()
     elif world_size > 1:
         # Using multiple gpus with multiple processes.
diff --git a/tests/configs/test_default.yml b/tests/configs/test_default.yml
index eba2cf917..8e3571895 100644
--- a/tests/configs/test_default.yml
+++ b/tests/configs/test_default.yml
@@ -34,6 +34,9 @@ data:
     data_directory: data/blender/lego_test
     dataset_type: blender
     downscale_factor: 4
+  
+  # validation options
+  validation_image_indices: [0]
 
   # dataloader options
   dataloader: