modularizing hydra config and setting defaults (#46)

* lint

* reconfigured modules in yaml

README.md

@@ -41,7 +41,7 @@ python scripts/run_train.py machine_config.num_gpus=1
 ```
 
 # Logging and profiling features
-We provide many logging functionalities for timing and/or tracking losses during training. All of these loggers are configurable via `configs/logging_configs.yml`
+We provide many logging functionalities for timing and/or tracking losses during training. All of these loggers are configurable via `configs/logging.yml`
 
 1. **Writer**: Logs losses and generated images during training to a specified output stream. Specify the type of writer (Tensorboard, local directory, Weights and Biases), and how often to log in the config.
 

configs/default.yaml

@@ -1,22 +1,11 @@
 defaults:
   # This makes the logging module logs colorful
+  - data: default_data.yaml
+  - graph: default_graph.yaml
+  - logging: default_logging.yaml
   - override hydra/job_logging: colorlog
   - override hydra/hydra_logging: colorlog
 
-logging_configs:
-  steps_per_log: 10
-  # stats tracker logs configurable training stats per iteration to terminal (e.g. rays/s, PSNR)
-  enable_stats: True
-  stats_tracker:
-    max_history: 10
-    stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME] # see mattport/utils/stats_tracker.py for options
-  # writer logs losses and images per iteration
-  writer: 
-    type: TensorboardWriter
-    save_dir: './'
-  # profiler logs run times of functions and prints at end of training
-  enable_profiler: True
-
 hydra:
   run:
     # Configure output dir of each experiment programmatically from the arguments
@@ -26,61 +15,9 @@ hydra:
 experiment_name: blender_lego
 method_name: vanilla_nerf
 
-# pretrained loading options
-resume_train:
-  load_dir: # pretrained model directory
-  load_step: # pretrained model's step number
-
-max_num_iterations: 1000000
-steps_per_save: null
-steps_per_test: 100
-
-model_dir: /tmp/mattport_models/
-
-# dataset options
-dataset:
-  data_directory: data/blender/lego
-  dataset_type: blender
-  downscale_factor: 4
-
-# dataloader options
-dataloader:
-  num_workers: 4
-  shuffle: true
-  num_images_to_sample_from: 8
-  num_rays_per_batch: 1024
-  num_times_to_repeat_images: 40
-
-# optimizer options
-param_groups:
-  fields:
-    optimizer:
-      type: Adam
-      lr: 1e-4
-    scheduler:
-      type: StepLR
-      step_size: 30 #  Period of learning rate decay.
-      gamma: 0.1 # Multiplicative factor of learning rate decay. Default: 0.1.
-  cameras:
-    optimizer:
-      type: SGD
-      lr: 1e-4
-    scheduler:
-      type: ExponentialLR
-      gamma: 0.1
-
 # machine options
 machine_config:
   num_gpus: 1
   num_machines: 1
   machine_rank: 0
   dist_url: auto
-
-# graph options
-graph:
-  # NOTE(ethan): currently setting to nothing b/c we're still designing the config format...
-  _target_: mattport.nerf.graph.vanilla_nerf.NeRFGraph
-  near_plane: 1.0
-  far_plane: 6.0
-  num_coarse_samples: 64
-  num_importance_samples: 64

configs/graph/default_graph.yaml

@@ -0,0 +1,36 @@
+# saving model specifications
+model_dir: /tmp/mattport_models/ 
+steps_per_save: null
+steps_per_test: 100
+max_num_iterations: 1000000
+
+# pretrained loading options
+resume_train:
+  load_dir: # pretrained model directory
+  load_step: # pretrained model's step number
+
+# NOTE(ethan): currently setting to nothing b/c we're still designing the config format...
+network:
+  _target_: mattport.nerf.graph.vanilla_nerf.NeRFGraph
+  near_plane: 1.0
+  far_plane: 6.0
+  num_coarse_samples: 64
+  num_importance_samples: 64
+
+# optimizer options for network
+param_groups:
+  fields:
+    optimizer:
+      type: Adam
+      lr: 1e-4
+    scheduler:
+      type: StepLR
+      step_size: 30 #  Period of learning rate decay.
+      gamma: 0.1 # Multiplicative factor of learning rate decay. Default: 0.1.
+  cameras:
+    optimizer:
+      type: SGD
+      lr: 1e-4
+    scheduler:
+      type: ExponentialLR
+      gamma: 0.1

configs/logging/default_logging.yaml

@@ -0,0 +1,12 @@
+steps_per_log: 10
+# stats tracker logs configurable training stats per iteration to terminal (e.g. rays/s, PSNR)
+enable_stats: True
+stats_tracker:
+  max_history: 10
+  stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME] # see mattport/utils/stats_tracker.py for options
+# writer logs losses and images per iteration
+writer: 
+  type: TensorboardWriter
+  save_dir: './'
+# profiler logs run times of functions and prints at end of training
+enable_profiler: True

mattport/nerf/trainer.py

@@ -44,17 +44,17 @@ class Trainer:
         self.start_step = 0
         # logging variables
         self.is_main_thread = local_rank % world_size == 0
-        writer = getattr(mattport.utils.writer, self.config.logging_configs.writer.type)
-        self.writer = writer(local_rank, world_size, save_dir=self.config.logging_configs.writer.save_dir)
+        writer = getattr(mattport.utils.writer, self.config.logging.writer.type)
+        self.writer = writer(local_rank, world_size, save_dir=self.config.logging.writer.save_dir)
         self.stats = StatsTracker(config, self.is_main_thread)
-        if not profiler.PROFILER and self.config.logging_configs.enable_profiler:
+        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)
 
     @profiler.time_function
     def setup(self):
         """Setup the Trainer by calling other setup functions."""
-        dataset_inputs_dict = get_dataset_inputs_dict(**self.config.dataset)
+        dataset_inputs_dict = get_dataset_inputs_dict(**self.config.data.dataset)
         self.setup_datasets(dataset_inputs_dict)
         self.setup_graph(dataset_inputs_dict["train"])
 
@@ -68,15 +68,15 @@ class Trainer:
         self.train_dataset = CollateIterDataset(
             self.train_image_dataset,
             collate_fn=lambda batch_list: collate_batch(
-                batch_list, self.config.dataloader.num_rays_per_batch, keep_full_image=False
+                batch_list, self.config.data.dataloader.num_rays_per_batch, keep_full_image=False
             ),
-            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,
             batch_size=1,
-            num_workers=self.config.dataloader.num_workers,
+            num_workers=self.config.data.dataloader.num_workers,
             collate_fn=collate_batch_size_one,
             pin_memory=True,
         )
@@ -93,12 +93,14 @@ 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(f"cuda:{self.local_rank}")
         self.setup_optimizers()
 
-        if self.config.resume_train.load_dir:
-            self.load_checkpoint(self.config.resume_train)
+        if self.config.graph.resume_train.load_dir:
+            self.load_checkpoint(self.config.graph.resume_train)
 
         if self.world_size > 1:
             self.graph = DDP(self.graph, device_ids=[self.local_rank])
@@ -107,7 +109,7 @@ class Trainer:
     def setup_optimizers(self):
         """_summary_"""
         # TODO(ethan): handle different world sizes
-        self.optimizers = Optimizers(self.config.param_groups, self.graph.get_param_groups())
+        self.optimizers = Optimizers(self.config.graph.param_groups, self.graph.get_param_groups())
 
     def load_checkpoint(self, load_config: DictConfig) -> int:
         """Load the checkpoint from the given path
@@ -165,7 +167,7 @@ class Trainer:
     def train(self) -> None:
         """_summary_"""
         train_start = time()
-        num_iterations = self.config.max_num_iterations
+        num_iterations = self.config.graph.max_num_iterations
         iter_dataset = iter(self.train_dataloader)
         for i, step in enumerate(range(self.start_step, self.start_step + num_iterations)):
             data_start = time()
@@ -179,12 +181,12 @@ class Trainer:
             )
             self.stats.update_time(Stats.ITER_TRAIN_TIME, iter_start, time(), step=step)
 
-            if step != 0 and step % self.config.logging_configs.steps_per_log == 0:
+            if step != 0 and step % self.config.logging.steps_per_log == 0:
                 self.writer.write_scalar_dict(loss_dict, step, group="Loss", prefix="train-")
                 # TODO: add the learning rates to tensorboard/logging
-            if step != 0 and self.config.steps_per_save and step % self.config.steps_per_save == 0:
-                self.save_checkpoint(self.config.model_dir, step)
-            if step % self.config.steps_per_test == 0:
+            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
             self.stats.print_stats(i / num_iterations)

mattport/utils/decorators.py

@@ -22,7 +22,7 @@ def check_stats_enabled(func: Callable) -> Callable:
 
     def wrapper(self, *args, **kwargs):
         ret = None
-        if self.config.logging_configs.enable_stats:
+        if self.config.logging.enable_stats:
             ret = func(self, *args, **kwargs)
         return ret
 
@@ -34,7 +34,7 @@ def check_profiler_enabled(func: Callable) -> Callable:
 
     def wrapper(self, *args, **kwargs):
         ret = None
-        if self.config.logging_configs.enable_stats:
+        if self.config.logging.enable_stats:
             ret = func(self, *args, **kwargs)
         return ret
 
@@ -59,9 +59,9 @@ def check_print_stats_step(func: Callable) -> Callable:
     def wrapper(self, *args, **kwargs):
         ret = None
         if (
-            self.step % self.config.logging_configs.steps_per_log == 0
-            or (self.config.steps_per_save and self.step % self.config.steps_per_save == 0)
-            or self.step % self.config.steps_per_test == 0
+            self.step % self.config.logging.steps_per_log == 0
+            or (self.config.graph.steps_per_save and self.step % self.config.graph.steps_per_save == 0)
+            or self.step % self.config.graph.steps_per_test == 0
         ):
             ret = func(self, *args, **kwargs)
         return ret

mattport/utils/profiler.py

@@ -37,7 +37,7 @@ class Profiler:
 
     def __init__(self, config: DictConfig, is_main_thread: bool):
         self.config = config
-        if self.config.logging_configs.enable_profiler:
+        if self.config.logging.enable_profiler:
             self.is_main_thread = is_main_thread
         self.profiler_dict = {}
 

mattport/utils/stats_tracker.py

@@ -30,15 +30,15 @@ class StatsTracker:
 
     def __init__(self, config: DictConfig, is_main_thread: bool):
         self.config = config
-        if self.config.logging_configs.enable_stats:
+        if self.config.logging.enable_stats:
             self.is_main_thread = is_main_thread
-            self.max_history = self.config.logging_configs.stats_tracker.max_history
+            self.max_history = self.config.logging.stats_tracker.max_history
             self.step = 0
             self.stats_dict = {}
             self.past_stats = []
             self.new_key = False
             self.stats_to_track = set()
-            for name in self.config.logging_configs.stats_tracker.stats_to_track:
+            for name in self.config.logging.stats_tracker.stats_to_track:
                 self.stats_to_track.add(Stats[name])
             logging.info("Successfully set up StatsTracker")
         else:
@@ -85,7 +85,7 @@ class StatsTracker:
 
             if name == Stats.ITER_TRAIN_TIME and Stats.ETA in self.stats_to_track:
                 # update ETA if logging iteration train time
-                remain_iter = self.config.max_num_iterations - step
+                remain_iter = self.config.graph.max_num_iterations - step
                 self.stats_dict[Stats.ETA] = remain_iter * self.stats_dict[name]
 
     @check_print_stats_step

scripts/run_train.py

@@ -176,7 +176,7 @@ def launch(
                 process.join()
                 logger.info("Process %s finished", str(i))
         finally:
-            if config.logging_configs.enable_profiler:
+            if config.logging.enable_profiler:
                 profiler.PROFILER.print_profile()
     else:
         try:
@@ -184,7 +184,7 @@ def launch(
         except KeyboardInterrupt:
             pass
         finally:
-            if config.logging_configs.enable_profiler:
+            if config.logging.enable_profiler:
                 profiler.PROFILER.print_profile()
 
 

tests/configs/test_default.yml

@@ -5,7 +5,7 @@ defaults:
   - override hydra/job_logging: colorlog
   - override hydra/hydra_logging: colorlog
 
-logging_configs:
+logging:
   steps_per_log: 10
   # stats tracker logs configurable training stats per iteration to terminal (e.g. rays/s, PSNR)
   enable_stats: True
@@ -28,48 +28,20 @@ hydra:
 experiment_name: blender_lego
 method_name: vanilla_nerf
 
-# pretrained loading options
-resume_train:
-  load_dir: # pretrained model directory
-  load_step: # pretrained model's step number
+data:
+  # dataset options
+  dataset:
+    data_directory: data/blender/lego_test
+    dataset_type: blender
+    downscale_factor: 4
 
-max_num_iterations: 10
-steps_per_save: null
-steps_per_test: 100
-
-model_dir: /tmp/mattport_models/
-
-# dataset options
-dataset:
-  data_directory: data/blender/lego_test
-  dataset_type: blender
-  downscale_factor: 4
-
-# dataloader options
-dataloader:
-  num_workers: 4
-  shuffle: true
-  num_images_to_sample_from: 1
-  num_rays_per_batch: 1024
-  num_times_to_repeat_images: 40
-
-# optimizer options
-param_groups:
-  fields:
-    optimizer:
-      type: Adam
-      lr: 1e-4
-    scheduler:
-      type: StepLR
-      step_size: 30 #  Period of learning rate decay.
-      gamma: 0.1 # Multiplicative factor of learning rate decay. Default: 0.1.
-  cameras:
-    optimizer:
-      type: SGD
-      lr: 1e-4
-    scheduler:
-      type: ExponentialLR
-      gamma: 0.1
+  # dataloader options
+  dataloader:
+    num_workers: 4
+    shuffle: true
+    num_images_to_sample_from: 1
+    num_rays_per_batch: 1024
+    num_times_to_repeat_images: 40
 
 # machine options
 machine_config:
@@ -80,9 +52,37 @@ machine_config:
 
 # graph options
 graph:
-  # NOTE(ethan): currently setting to nothing b/c we're still designing the config format...
-  _target_: mattport.nerf.graph.vanilla_nerf.NeRFGraph
-  near_plane: 1.0
-  far_plane: 6.0
-  num_coarse_samples: 64
-  num_importance_samples: 64
+  model_dir: /tmp/mattport_models/
+  max_num_iterations: 10
+  steps_per_save: null
+  steps_per_test: 100
+
+  # pretrained loading options
+  resume_train:
+    load_dir: # pretrained model directory
+    load_step: # pretrained model's step number
+
+  network:
+    # NOTE(ethan): currently setting to nothing b/c we're still designing the config format...
+    _target_: mattport.nerf.graph.vanilla_nerf.NeRFGraph
+    near_plane: 1.0
+    far_plane: 6.0
+    num_coarse_samples: 64
+    num_importance_samples: 64
+  # optimizer options
+  param_groups:
+    fields:
+      optimizer:
+        type: Adam
+        lr: 1e-4
+      scheduler:
+        type: StepLR
+        step_size: 30 #  Period of learning rate decay.
+        gamma: 0.1 # Multiplicative factor of learning rate decay. Default: 0.1.
+    cameras:
+      optimizer:
+        type: SGD
+        lr: 1e-4
+      scheduler:
+        type: ExponentialLR
+        gamma: 0.1