Printing profile on sig kill. (#44)

* added skeleton for terminal logger

* done skeleton terminal writer logic... todo: actual calls to the write classes

* moved lambda to top function

* added logic to support timing

* prettier printing

* prettier printing

* lint

* separated out timing and other stats variables. printing in dict format

* keeping logging history and wiping rest

* updated logic to calculate total times instead of average times

* made stats tracker configurable

* moved stats tracker and profiler to own class file; implemented basic functionality for profiler

* lint removing duplicate function

* update readme; moved around config variables

* profiler prints on sig kill

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

@@ -47,7 +47,8 @@ class Trainer:
         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)
+        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)
 
     @profiler.time_function
@@ -178,7 +179,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:
@@ -190,7 +191,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)
 
@@ -174,8 +175,17 @@ def launch(
                     process.terminate()
                 process.join()
                 logger.info("Process %s finished", str(i))
+        finally:
+            if config.logging_configs.enable_profiler:
+                profiler.PROFILER.print_profile()
     else:
-        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()
 
 
 @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/'
+    save_dir: './'
+  # 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/