From bbfe815504baf9b46c9f3b74c5c8e31b7fc22048 Mon Sep 17 00:00:00 2001 From: evonneng Date: Fri, 20 May 2022 18:38:18 +0000 Subject: [PATCH] added logger ability to log everything without erasing --- configs/logging/default_logging.yaml | 5 +- mattport/nerf/graph/nerfw.py | 2 +- mattport/nerf/graph/vanilla_nerf.py | 2 +- mattport/nerf/trainer.py | 14 +++-- mattport/utils/writer.py | 85 +++++++++++++++++++--------- 5 files changed, 72 insertions(+), 36 deletions(-) diff --git a/configs/logging/default_logging.yaml b/configs/logging/default_logging.yaml index 8938d1233..9dfaf40e1 100644 --- a/configs/logging/default_logging.yaml +++ b/configs/logging/default_logging.yaml @@ -1,12 +1,13 @@ steps_per_log: 10 -max_history: 20 +max_buffer_size: 20 # writer logs losses and images per iteration writer: TensorboardWriter: save_dir: './' LocalWriter: save_dir: './' - stats_to_track: ["Data Load (time)", "Train Iter (time)", "Rays Per Sec (time)", "Test PSNR"] + stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME, RAYS_PER_SEC, CURR_TEST_PSNR] + max_log_size: 20 # if 0, logs everything with no erasing # profiler logs run times of functions and prints at end of training enable_profiler: True \ No newline at end of file diff --git a/mattport/nerf/graph/nerfw.py b/mattport/nerf/graph/nerfw.py index b31061266..0b1a09b1a 100644 --- a/mattport/nerf/graph/nerfw.py +++ b/mattport/nerf/graph/nerfw.py @@ -212,4 +212,4 @@ class NerfWGraph(NeRFGraph): row2 = torch.cat([depth_fine, depth_fine_static, depth_coarse], dim=-2) combined_image = torch.cat([row0, row1, row2], dim=-3) - writer.put_scalar(name="image_idx_{image_idx}-nerfw", image=combined_image, step=step, group="img") + writer.put_image(name="image_idx_{image_idx}-nerfw", image=combined_image, step=step, group="img") diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index f1ba247ba..a4b41e9e7 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -251,6 +251,6 @@ class NeRFGraph(Graph): writer.put_scalar(name=f"val_{image_idx}", scalar=float(fine_ssim), step=step, group="ssim") writer.put_scalar(name=f"val_{image_idx}", scalar=float(fine_lpips), step=step, group="lpips") - writer.put_scalar(name="Test PSNR", scalar=float(fine_psnr), step=step) + writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(fine_psnr), step=step) return fine_psnr.item() diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index e6f654a24..e86be776d 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -202,13 +202,17 @@ class Trainer: data_start = time() batch = next(iter_dataset) writer.put_time( - name="Data Load (time)", start_time=data_start, end_time=time(), step=step, avg_over_iters=True + name=writer.EventName.ITER_LOAD_TIME, + start_time=data_start, + end_time=time(), + step=step, + avg_over_iters=True, ) iter_start = time() loss_dict = self.train_iteration(batch, step) writer.put_time( - name="Rays Per Sec (time)", + name=writer.EventName.RAYS_PER_SEC, start_time=iter_start, end_time=time(), step=step, @@ -216,7 +220,7 @@ class Trainer: avg_over_batch=batch["indices"].shape[0], ) writer.put_time( - name="Train Iter (time)", + name=writer.EventName.ITER_TRAIN_TIME, start_time=iter_start, end_time=time(), step=step, @@ -234,7 +238,9 @@ class Trainer: _ = self.test_image(image_idx=image_idx, step=step) self._write_out_storage(step) - writer.put_time(name="Train Total (time)", start_time=train_start, end_time=time(), step=num_iterations) + writer.put_time( + name=writer.EventName.TOTAL_TRAIN_TIME, start_time=train_start, end_time=time(), step=num_iterations + ) self._write_out_storage(num_iterations) def _write_out_storage(self, step): diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index ecacd7567..47130761a 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -25,12 +25,20 @@ EVENT_STORAGE = [] GLOBAL_BUFFER = {} -class EventType(enum.Enum): - """Possible Event types that can be logged +class EventName(enum.Enum): + """Names of possible events that can be logged via Local Writer for convenience. + see config/logging/default_logging.yaml""" - Args: - enum (_type_): _description_ - """ + ITER_LOAD_TIME = "Data Load (time)" + ITER_TRAIN_TIME = "Train Iter (time)" + RAYS_PER_SEC = "Rays Per Sec (time)" + TOTAL_TRAIN_TIME = "Train Total (time)" + ETA = "ETA (time)" + CURR_TEST_PSNR = "Test PSNR" + + +class EventType(enum.Enum): + """Possible Event types and their associated write function""" IMAGE = "write_image" SCALAR = "write_scalar" @@ -40,6 +48,9 @@ class EventType(enum.Enum): @check_main_thread def put_image(name, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None): """Setter function to place images into the queue to be written out""" + if isinstance(name, EventName): + name = name.value + EVENT_STORAGE.append( {"name": name, "write_type": EventType.IMAGE, "event": image, "step": step, "group": group, "prefix": prefix} ) @@ -48,6 +59,9 @@ def put_image(name, image: TensorType["H", "W", "C"], step: int, group: str = No @check_main_thread def put_scalar(name: str, scalar: float, step: int, group: str = None, prefix: str = None): """Setter function to place scalars into the queue to be written out""" + if isinstance(name, EventName): + name = name.value + GLOBAL_BUFFER["new_key"] = not name in GLOBAL_BUFFER["events"] or GLOBAL_BUFFER["new_key"] GLOBAL_BUFFER["events"][name] = scalar EVENT_STORAGE.append( @@ -89,6 +103,9 @@ def put_time( avg_over_batch (int): if set, the size of the batch for which we take the average over (batch/second) update_eta (bool): if True, update the ETA. should only be set for the training iterations/s """ + if isinstance(name, EventName): + name = name.value + GLOBAL_BUFFER["step"] = step GLOBAL_BUFFER["new_key"] = not name in GLOBAL_BUFFER["events"] or GLOBAL_BUFFER["new_key"] val = end_time - start_time @@ -99,7 +116,7 @@ def put_time( curr_event = GLOBAL_BUFFER["events"].get(name, {"buffer": [], "avg": 0}) curr_buffer = curr_event["buffer"] curr_avg = curr_event["avg"] - if len(curr_buffer) >= GLOBAL_BUFFER["max_history"]: + if len(curr_buffer) >= GLOBAL_BUFFER["max_buffer_size"]: curr_buffer.pop(0) curr_buffer.append(val) curr_avg = sum(curr_buffer) / len(curr_buffer) @@ -114,7 +131,7 @@ def put_time( remain_iter = GLOBAL_BUFFER["max_iter"] - step remain_time = remain_iter * GLOBAL_BUFFER["events"][name]["avg"] GLOBAL_BUFFER["events"]["ETA"] = remain_time - put_scalar("ETA", remain_time, step, group, prefix) + put_scalar(EventName.ETA, remain_time, step, group, prefix) @check_main_thread @@ -140,16 +157,16 @@ def setup_event_writers(config: DictConfig) -> None: logging_configs = config.logging.writer for writer_type in logging_configs: writer_class = getattr(mattport.utils.writer, writer_type) - curr_config = logging_configs[writer_type] + writer_config = logging_configs[writer_type] if writer_type == "LocalWriter": - curr_writer = writer_class(curr_config.save_dir, curr_config.stats_to_track, config) + curr_writer = writer_class(writer_config.save_dir, writer_config.stats_to_track, writer_config.max_log_size) else: - curr_writer = writer_class(curr_config.save_dir) + curr_writer = writer_class(writer_config.save_dir) EVENT_WRITERS.append(curr_writer) ## configure all the global buffer basic information GLOBAL_BUFFER["max_iter"] = config.graph.max_num_iterations - GLOBAL_BUFFER["max_history"] = config.logging.max_history + GLOBAL_BUFFER["max_buffer_size"] = config.logging.max_buffer_size GLOBAL_BUFFER["steps_per_log"] = config.logging.steps_per_log GLOBAL_BUFFER["new_key"] = True GLOBAL_BUFFER["events"] = {} @@ -278,12 +295,17 @@ def _consolidate_events(): class LocalWriter(Writer): """Local Writer Class""" - def __init__(self, save_dir: str, stats_to_track: ListConfig, config: DictConfig): + def __init__(self, save_dir: str, stats_to_track: ListConfig, max_log_size: int = 0): + """ + Args: + stats_to_track (ListConfig): the names of stats that should be logged. + max_log size (int): max number of lines that will be logged to teminal. + """ super().__init__(save_dir) - self.stats_to_track = stats_to_track + self.stats_to_track = [EventName[name].value for name in stats_to_track] + self.max_log_size = max_log_size self.past_stats = [] self.max_mssg_len = 0 - self.config = config def write_image( self, name: str, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None @@ -301,18 +323,22 @@ class LocalWriter(Writer): def handle_header(self, latest_map): """helper to handle the printing of the header labels""" - if len(self.past_stats) == 0 or GLOBAL_BUFFER["new_key"]: + if (not self.max_log_size and self.max_mssg_len == 0) or ( + self.max_log_size and (len(self.past_stats) == 0 or GLOBAL_BUFFER["new_key"]) + ): mssg = f"{'Step (% Done)':<20}" for name, _ in latest_map.items(): if name in self.stats_to_track: mssg += f"{name:<20} " - if len(self.past_stats) > 1: - _cursorup(len(self.past_stats) - 1) + if self.max_log_size: + if len(self.past_stats) > 1: + _cursorup(len(self.past_stats) - 1) print(mssg) print("-" * len(mssg)) - if len(self.past_stats) > 0: - for mssg in self.past_stats: - print(mssg) + if self.max_log_size: + if len(self.past_stats) > 0: + for mssg in self.past_stats: + print(mssg) self.max_mssg_len = len(mssg) def handle_stats(self, latest_map, padding=" "): @@ -330,12 +356,15 @@ class LocalWriter(Writer): curr_mssg += f"{v:<20} " # update the history buffer - if len(self.past_stats) > GLOBAL_BUFFER["max_history"]: - self.past_stats.pop(0) - _cursorup(len(self.past_stats)) - elif len(self.past_stats) > 1: - _cursorup(len(self.past_stats) - 1) + if self.max_log_size: + if len(self.past_stats) > self.max_log_size: + self.past_stats.pop(0) + _cursorup(len(self.past_stats)) + elif len(self.past_stats) > 1: + _cursorup(len(self.past_stats) - 1) - for mssg in self.past_stats: - print(f"{mssg:{padding}<{self.max_mssg_len}}") - self.past_stats.append(curr_mssg) + for mssg in self.past_stats: + print(f"{mssg:{padding}<{self.max_mssg_len}}") + self.past_stats.append(curr_mssg) + else: + print(curr_mssg)