merged master into this

README.md

@@ -1,7 +1,8 @@
 # MattPort
 
 
-# Setup the environment
+# Quickstart
+## Installation: Setup the environment
 
 ```
 # Clone the repo
@@ -19,7 +20,7 @@ python setup.py develop
 pytest tests
 ```
 
-# Getting the data
+## Getting the data
 
 Download the original NeRF dataset and put it in the following format.
 
@@ -30,7 +31,7 @@ Download the original NeRF dataset and put it in the following format.
         - lego/
 ```
 
-# Training a model
+## Training a model
 
 ```
 # Run with the default config
@@ -39,6 +40,16 @@ python scripts/run_train.py
 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`
+
+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.
+
+2. **StatsTracker**: Computes select statistics and prints to the terminal for local debugging. Specify the kinds of statistics to be logged, how often to log, and the maximum history buffer size in the config.
+
+3. **Profiler**: Computes the average total time of execution for any function with the `@profiler.time_function` decorator. Prints out the full profile at the termination of training or the program.
+
+
 # Setting up Jupyter
 
 ```

configs/default.yaml

@@ -4,17 +4,18 @@ defaults:
   - 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
-  # 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: "./"
+  # profiler logs run times of functions and prints at end of training
+  enable_profiler: True
 
 hydra:
   run:
@@ -32,7 +33,6 @@ resume_train:
 
 max_num_iterations: 1000000
 steps_per_save: null
-steps_per_log: 10
 steps_per_test: 100
 
 model_dir: /tmp/mattport_models/

mattport/nerf/trainer.py

@@ -55,7 +55,8 @@ class Trainer:
         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)
         self.stats = StatsTracker(config, self.is_main_thread)
-        profiler.PROFILER = profiler.Profiler(config, self.is_main_thread)
+        if not profiler.PROFILER and self.config.logging_configs.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"
 
@@ -193,7 +194,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:
+            if step != 0 and step % self.config.logging_configs.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:
@@ -205,7 +206,6 @@ class Trainer:
 
         self.stats.update_time(Stats.TOTAL_TRAIN_TIME, train_start, time(), step=-1)
         self.stats.print_stats(-1)
-        profiler.PROFILER.print_profile()
 
     @profiler.time_function
     def train_iteration(self, batch: dict, step: int):

mattport/utils/decorators.py

@@ -59,7 +59,7 @@ def check_print_stats_step(func: Callable) -> Callable:
     def wrapper(self, *args, **kwargs):
         ret = None
         if (
-            self.step % self.config.steps_per_log == 0
+            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
         ):

mattport/utils/profiler.py

@@ -1,8 +1,8 @@
 """
 Profiler base class and functionality
 """
-import sys
 import logging
+import sys
 import time
 from typing import Callable
 
@@ -19,18 +19,19 @@ def time_function(func: Callable) -> Callable:
     def wrapper(*args, **kwargs):
         start = time.time()
         ret = func(*args, **kwargs)
-        vals = vars(sys.modules[func.__module__])
-        class_str = ""
-        for attr in func.__qualname__.split(".")[:-1]:
-            class_str += f"{vals[attr].__qualname__}_"
-        class_str += func.__name__
-        PROFILER.update_time(class_str, start, time.time())
+        if PROFILER:
+            vals = vars(sys.modules[func.__module__])
+            class_str = ""
+            for attr in func.__qualname__.split(".")[:-1]:
+                class_str += f"{vals[attr].__qualname__}_"
+            class_str += func.__name__
+            PROFILER.update_time(class_str, start, time.time())
         return ret
 
     return wrapper
 
 
-@decorate_all([check_main_thread, check_profiler_enabled])
+@decorate_all([check_profiler_enabled, check_main_thread])
 class Profiler:
     """Profiler class"""
 

mattport/utils/stats_tracker.py

@@ -24,7 +24,7 @@ class Stats(enum.Enum):
     ETA = "ETA (time)"
 
 
-@decorate_all([check_main_thread, check_stats_enabled])
+@decorate_all([check_stats_enabled, check_main_thread])
 class StatsTracker:
     """Stats Tracker class"""
 

scripts/run_train.py

@@ -16,6 +16,7 @@ import torch.multiprocessing as mp
 from omegaconf import DictConfig
 
 from mattport.nerf.trainer import Trainer
+from mattport.utils import profiler
 
 logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG)
 
@@ -147,7 +148,13 @@ def launch(
         main_func(local_rank=0, world_size=0, config=config, cpu=True)
     elif world_size == 1:
         # Using one gpu and one process.
-        main_func(local_rank=0, world_size=1, config=config)
+        try:
+            main_func(local_rank=0, world_size=1, config=config)
+        except KeyboardInterrupt:
+            pass
+        finally:
+            if config.logging_configs.enable_profiler:
+                profiler.PROFILER.print_profile()
     elif world_size > 1:
         # Using multiple gpus with multiple processes.
         if dist_url == "auto":
@@ -182,6 +189,9 @@ def launch(
                     process.terminate()
                 process.join()
                 logger.info("Process %s finished", str(i))
+        finally:
+            if config.logging_configs.enable_profiler:
+                profiler.PROFILER.print_profile()
 
 
 @hydra.main(config_path="../configs", config_name="default.yaml")

tests/configs/test_default.yml

@@ -6,17 +6,18 @@ defaults:
   - 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
-  # 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/"
+  # profiler logs run times of functions and prints at end of training
+  enable_profiler: True
 
 hydra:
   run:
@@ -34,7 +35,6 @@ resume_train:
 
 max_num_iterations: 10
 steps_per_save: null
-steps_per_log: 10
 steps_per_test: 100
 
 model_dir: /tmp/mattport_models/