added general testing, fixed up config, made writer general

configs/default.yaml

@@ -3,12 +3,18 @@ defaults:
   - override hydra/job_logging: colorlog
   - override hydra/hydra_logging: colorlog
 
-debug:
+logging_configs:
+  # 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
+  # profiler logs run times of functions and prints at end of training
   enable_profiler: True
+  # writer logs losses and images per iteration
+  writer: 
+    type: TensorboardWriter
+    save_dir: './'
 
 hydra:
   run:

mattport/nerf/trainer.py

@@ -13,6 +13,7 @@ from omegaconf import DictConfig
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import DataLoader
 
+import mattport.utils.writer
 from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one
 from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch
 from mattport.nerf.dataset.utils import get_dataset_inputs
@@ -21,7 +22,6 @@ from mattport.nerf.optimizers import Optimizers
 from mattport.utils import profiler
 from mattport.utils.decorators import check_main_thread
 from mattport.utils.stats_tracker import Stats, StatsTracker
-from mattport.utils.writer import TensorboardWriter
 
 logging.getLogger("PIL").setLevel(logging.WARNING)
 
@@ -45,9 +45,8 @@ class Trainer:
         self.start_step = 0
         # logging variables
         self.is_main_thread = local_rank % world_size == 0
-        self.tensorboard_writer = TensorboardWriter(
-            local_rank, world_size, save_dir=os.path.join(os.getcwd(), "writer")
-        )
+        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)
         self.stats = StatsTracker(config, self.is_main_thread)
         profiler.PROFILER = profiler.Profiler(config, self.is_main_thread)
 
@@ -167,7 +166,7 @@ class Trainer:
             self.stats.update_time(Stats.ITER_TRAIN_TIME, iter_start, time(), step=step)
 
             if step != 0 and step % self.config.steps_per_log == 0:
-                self.tensorboard_writer.write_scalar_dict(loss_dict, step, group="Loss", prefix="train-")
+                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)
@@ -220,10 +219,8 @@ class Trainer:
             rgb_fine = torch.cat(rgb_fine).view(image_height, image_width, 3).detach().cpu()
 
         combined_image = torch.cat([image, rgb_coarse, rgb_fine], dim=1)
-        self.tensorboard_writer.write_image(
-            f"image_idx_{image_idx}-rgb_coarse_fine", combined_image, step, group="test"
-        )
+        self.writer.write_image(f"image_idx_{image_idx}-rgb_coarse_fine", combined_image, step, group="test")
 
         fine_psnr = get_psnr(image, rgb_fine)
         self.stats.update_value(Stats.CURR_TEST_PSNR, fine_psnr, step)
-        self.tensorboard_writer.write_scalar(f"image_idx_{image_idx}-fine_psnr", fine_psnr, step, group="test")
+        self.writer.write_scalar(f"image_idx_{image_idx}-fine_psnr", fine_psnr, step, group="test")

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.debug.enable_stats:
+        if self.config.logging_configs.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.debug.enable_stats:
+        if self.config.logging_configs.enable_stats:
             ret = func(self, *args, **kwargs)
         return ret
 

mattport/utils/profiler.py

@@ -36,7 +36,7 @@ class Profiler:
 
     def __init__(self, config: DictConfig, is_main_thread: bool):
         self.config = config
-        if self.config.debug.enable_profiler:
+        if self.config.logging_configs.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.debug.enable_stats:
+        if self.config.logging_configs.enable_stats:
             self.is_main_thread = is_main_thread
-            self.max_history = self.config.debug.stats_tracker.max_history
+            self.max_history = self.config.logging_configs.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.debug.stats_tracker.stats_to_track:
+            for name in self.config.logging_configs.stats_tracker.stats_to_track:
                 self.stats_to_track.add(Stats[name])
             logging.info("Successfully set up StatsTracker")
         else:

tests/configs/test_default.yml

@@ -0,0 +1,86 @@
+defaults:
+  # This makes the logging module logs colorful
+  - override hydra/job_logging: colorlog
+  - override hydra/hydra_logging: colorlog
+
+logging_configs:
+  # 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
+  # profiler logs run times of functions and prints at end of training
+  enable_profiler: True
+  # writer logs losses and images per iteration
+  writer: 
+    type: TensorboardWriter
+    save_dir: '/tmp/'
+
+hydra:
+  run:
+    # Configure output dir of each experiment programmatically from the arguments
+    # Example "outputs/mnist/classifier/baseline/2021-03-10-141516"
+    dir: outputs/${experiment_name}/${method_name}/${now:%Y-%m-%d_%H%M%S}
+
+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: 10
+steps_per_save: null
+steps_per_log: 10
+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

tests/test_run_train.py

@@ -0,0 +1,16 @@
+"""
+Default test to make sure train runs
+"""
+import pytest
+import omegaconf
+from mattport.utils.io import get_absolute_path
+
+from scripts.run_train import main
+
+
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_run_train():
+    """test run train script works properly"""
+    test_config = get_absolute_path("./tests/configs/test_default.yml")
+    cfg = omegaconf.OmegaConf.load(test_config)
+    main(cfg)