From 59eeaca93b4223ab5564ebe296db45ecea895342 Mon Sep 17 00:00:00 2001 From: evonneng Date: Fri, 20 May 2022 02:19:44 +0000 Subject: [PATCH 1/4] removed stats tracker; consolidate writers into local writer --- configs/logging/default_logging.yaml | 12 +- mattport/nerf/graph/nerfw.py | 2 +- mattport/nerf/graph/semantic_nerf.py | 2 +- mattport/nerf/graph/vanilla_nerf.py | 20 +-- mattport/nerf/optimizers.py | 2 +- mattport/nerf/trainer.py | 60 ++++--- mattport/utils/decorators.py | 34 +--- mattport/utils/stats_tracker.py | 192 --------------------- mattport/utils/writer.py | 240 ++++++++++++++++++++++----- 9 files changed, 248 insertions(+), 316 deletions(-) delete mode 100644 mattport/utils/stats_tracker.py diff --git a/configs/logging/default_logging.yaml b/configs/logging/default_logging.yaml index 59ef1333b..a1279499c 100644 --- a/configs/logging/default_logging.yaml +++ b/configs/logging/default_logging.yaml @@ -1,12 +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, RAYS_PER_SEC, CURR_TEST_PSNR] # see mattport/utils/stats_tracker.py for options # writer logs losses and images per iteration writer: - TensorboardWriter: + # TensorboardWriter: + # save_dir: './' + LocalWriter: save_dir: './' + max_history: 20 + stats_to_track: ["Data Load (time)", "Train Iter (time)", "Rays Per Sec (time)", "Test PSNR"] + # 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 06ae5d81f..b31061266 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.write_image(name="image_idx_{image_idx}-nerfw", image=combined_image, step=step, group="img") + writer.put_scalar(name="image_idx_{image_idx}-nerfw", image=combined_image, step=step, group="img") diff --git a/mattport/nerf/graph/semantic_nerf.py b/mattport/nerf/graph/semantic_nerf.py index 02357caf6..726c54d77 100644 --- a/mattport/nerf/graph/semantic_nerf.py +++ b/mattport/nerf/graph/semantic_nerf.py @@ -155,4 +155,4 @@ class SemanticNerfGraph(NeRFGraph): self.stuff_classes ) semantic_labels_image = visualization.apply_colormap(semantic_labels[..., None]) - writer.write_image(name="image_idx_{image_idx}", image=semantic_labels_image, step=step, group="semantics") + writer.put_image(name="image_idx_{image_idx}", image=semantic_labels_image, step=step, group="semantics") diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index ffb90bea2..5b30d940a 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -23,7 +23,7 @@ from mattport.nerf.sampler import PDFSampler, UniformSampler from mattport.structures import colors from mattport.structures.rays import PointSamples, RayBundle from mattport.nerf.colliders import NearFarCollider -from mattport.utils import stats_tracker, visualization, writer +from mattport.utils import visualization, writer from mattport.utils.misc import is_not_none @@ -231,9 +231,9 @@ class NeRFGraph(Graph): combined_acc = torch.cat([acc_coarse, acc_fine], dim=1) combined_depth = torch.cat([depth_coarse, depth_fine], dim=1) - writer.write_image(name=f"image_idx_{image_idx}", image=combined_rgb, step=step, group="img") - writer.write_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation") - writer.write_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth") + writer.put_image(name=f"image_idx_{image_idx}", image=combined_rgb, step=step, group="img") + writer.put_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation") + writer.put_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth") # Switch images from [H, W, C] to [1, C, H, W] for metrics computations image = torch.moveaxis(image, -1, 0)[None, ...] @@ -245,13 +245,11 @@ class NeRFGraph(Graph): fine_ssim = self.ssim(image, rgb_fine) fine_lpips = self.lpips(image, rgb_fine) - writer.write_scalar(name=f"val_{image_idx}-coarse", scalar=float(coarse_psnr), step=step, group="psnr") - writer.write_scalar(name=f"val_{image_idx}-fine", scalar=float(fine_psnr), step=step, group="psnr") - writer.write_scalar(name=f"val_{image_idx}", scalar=float(fine_ssim), step=step, group="ssim") - writer.write_scalar(name=f"val_{image_idx}", scalar=float(fine_lpips), step=step, group="lpips") + writer.put_scalar(name=f"val_{image_idx}-coarse", scalar=float(coarse_psnr), step=step, group="psnr") + writer.put_scalar(name=f"val_{image_idx}-fine", scalar=float(fine_psnr), step=step, group="psnr") + 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") - stats_tracker.update_stats( - {"name": stats_tracker.Stats.CURR_TEST_PSNR, "value": float(fine_psnr), "step": step} - ) + writer.put_scalar(name="Test PSNR", scalar=float(fine_psnr), step=step) return fine_psnr.item() diff --git a/mattport/nerf/optimizers.py b/mattport/nerf/optimizers.py index 917eeed0d..181ac1517 100644 --- a/mattport/nerf/optimizers.py +++ b/mattport/nerf/optimizers.py @@ -83,4 +83,4 @@ class Optimizers: scheduler.step() # TODO(ethan): clean this up. why is there indexing into a list? lr = scheduler.get_last_lr()[0] - writer.write_scalar(name=f"{param_group_name}", scalar=lr, step=step, group="learning_rate") + writer.put_scalar(name=f"{param_group_name}", scalar=lr, step=step, group="learning_rate") diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index ec2486a92..ee2232ba4 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -18,7 +18,7 @@ from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch from mattport.nerf.dataset.utils import DatasetInputs, get_dataset_inputs from mattport.nerf.optimizers import Optimizers -from mattport.utils import profiler, stats_tracker, writer +from mattport.utils import profiler, writer from mattport.utils.decorators import check_main_thread logging.getLogger("PIL").setLevel(logging.WARNING) @@ -51,8 +51,7 @@ class Trainer: self.optimizers = None self.start_step = 0 # logging variables - stats_tracker.setup_stats_tracker(config) - writer.setup_event_writers(config.logging.writer) + writer.setup_event_writers(config) profiler.setup_profiler(config.logging) self.device = "cpu" if self.world_size == 0 else f"cuda:{self.local_rank}" @@ -199,47 +198,54 @@ class Trainer: train_start = time() 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)): + for step in range(self.start_step, self.start_step + num_iterations): + self.step = step data_start = time() batch = next(iter_dataset) - stats_tracker.update_stats( - {"name": stats_tracker.Stats.ITER_LOAD_TIME, "start_time": data_start, "end_time": time(), "step": step} + writer.put_time( + name="Data 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) - stats_tracker.update_stats( - { - "name": stats_tracker.Stats.RAYS_PER_SEC, - "start_time": iter_start, - "end_time": time(), - "step": step, - "batch_size": batch["indices"].shape[0], - }, + writer.put_time( + name="Rays Per Sec (time)", + start_time=iter_start, + end_time=time(), + step=step, + avg_over_iters=True, + avg_over_batch=batch["indices"].shape[0], ) - stats_tracker.update_stats( - { - "name": stats_tracker.Stats.ITER_TRAIN_TIME, - "start_time": iter_start, - "end_time": time(), - "step": step, - } + writer.put_time( + name="Train Iter (time)", + start_time=iter_start, + end_time=time(), + step=step, + avg_over_iters=True, + update_eta=True, ) if step != 0 and step % self.config.logging.steps_per_log == 0: - writer.write_scalar_dict(scalar_dict=loss_dict, step=step, group="Loss", prefix="train-") + writer.put_dict(name="loss_dict", scalar_dict=loss_dict, step=step, group="Loss", prefix="train-") # TODO: add the learning rates to tensorboard/logging 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: # NOTE(ethan): we should still run this in dry-run mode! for image_idx in self.config.data.val_image_indices: _ = self.test_image(image_idx=image_idx, step=step) - stats_tracker.print_stats(i / num_iterations) + self._write_out_storage() - stats_tracker.update_stats( - {"name": stats_tracker.Stats.TOTAL_TRAIN_TIME, "start_time": train_start, "end_time": time()} - ) - stats_tracker.print_stats(1.0) + writer.put_time(name="Train Total (time)", start_time=train_start, end_time=time(), step=num_iterations) + self._write_out_storage() + + def _write_out_storage(self): + """Perform writes only during appropriate time steps""" + if ( + 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 + ): + writer.write_out_storage() @profiler.time_function def train_iteration(self, batch: dict, step: int): diff --git a/mattport/utils/decorators.py b/mattport/utils/decorators.py index f7b79f7bf..9d99c68f9 100644 --- a/mattport/utils/decorators.py +++ b/mattport/utils/decorators.py @@ -19,24 +19,12 @@ def decorate_all(decorators: List[Callable]) -> Callable: return decorate -def check_stats_enabled(func: Callable) -> Callable: - """Decorator: check if stats tracker is enabled""" - - def wrapper(self, *args, **kwargs): - ret = None - if self.config.logging.enable_stats: - ret = func(self, *args, **kwargs) - return ret - - return wrapper - - def check_profiler_enabled(func: Callable) -> Callable: """Decorator: check if profiler is enabled""" def wrapper(self, *args, **kwargs): ret = None - if self.config.enable_stats: + if self.config.enable_profiler: ret = func(self, *args, **kwargs) return ret @@ -46,26 +34,10 @@ def check_profiler_enabled(func: Callable) -> Callable: def check_main_thread(func: Callable) -> Callable: """Decorator: check if you are on main thread""" - def wrapper(self, *args, **kwargs): + def wrapper(*args, **kwargs): ret = None if comms.is_main_process(): - ret = func(self, *args, **kwargs) - return ret - - return wrapper - - -def check_print_stats_step(func: Callable) -> Callable: - """Decorator: check if it is time to print stats update""" - - def wrapper(self, *args, **kwargs): - ret = None - if ( - 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) + ret = func(*args, **kwargs) return ret return wrapper diff --git a/mattport/utils/stats_tracker.py b/mattport/utils/stats_tracker.py deleted file mode 100644 index 64f0fc5f5..000000000 --- a/mattport/utils/stats_tracker.py +++ /dev/null @@ -1,192 +0,0 @@ -""" -Stats tracker base class and functionality -""" -import datetime -import enum -import logging -from typing import Any, Dict - -from omegaconf import DictConfig - -from mattport.utils import comms -from mattport.utils.decorators import check_main_thread, check_print_stats_step, check_stats_enabled, decorate_all - -STATS_TRACKER = [] - - -def update_stats(args: Dict[str, Any]): - """update global stats tracker based on arguments""" - assert comms.is_main_process(), "Writing out with process other than main" - if "value" in args: - STATS_TRACKER[0].update_value(**args) - elif "start_time" in args: - STATS_TRACKER[0].update_time(**args) - else: - raise NotImplementedError - - -def print_stats(fraction_done: float): - """print current statistics information""" - assert comms.is_main_process(), "Writing out with process other than main" - STATS_TRACKER[0].print_stats(fraction_done) - - -def setup_stats_tracker(config: DictConfig): - """Initialization of stats tracker""" - if comms.is_main_process(): - STATS_TRACKER.append(StatsTracker(config)) - - -class Stats(enum.Enum): - """Possible Stats values for StatsTracker - - Args: - enum (_type_): _description_ - """ - - ITER_LOAD_TIME = "Data Load (ms)" - ITER_TRAIN_TIME = "Train Iter (ms)" - TOTAL_TRAIN_TIME = "Train Total (time)" - RAYS_PER_SEC = "Rays Per Sec (1/s)" - CURR_TEST_PSNR = "Test PSNR" - ETA = "ETA (time)" - - -@decorate_all([check_stats_enabled, check_main_thread]) -class StatsTracker: - """Stats Tracker class""" - - def __init__(self, config: DictConfig): - self.config = config - if self.config.logging.enable_stats: - 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.stats_tracker.stats_to_track: - self.stats_to_track.add(Stats[name]) - logging.info("Successfully set up StatsTracker") - else: - logging.info("StatsTracker disabled; enable in config.") - - def update_value(self, name: enum.Enum, value: float, step: int): - """update stats dictionary with key value pair - - Args: - name (enum.Enum): Enum name of statistic we are logging - value (float): value to update. - step (int): number of total iteration steps. - """ - if name in self.stats_to_track: - self.step = step - self.new_key = not name in self.stats_dict or self.new_key - self.stats_dict[name] = value - - def update_time( - self, name: enum.Enum, start_time: float, end_time: float, step: int = None, batch_size: int = None - ): - """update the stats dictionary with running averages/cumulative durations - - Args: - name (enum.Enum): Enum name of statistic we are logging - start_time (float): start time for the call in seconds - end_time (float): end time when the call finished executing in seconds - step (int): number of total iteration steps. Defaults to None. - if None, reports duration without averaging - batch_size (int, optional): total number of rays in a batch; - if None, reports duration instead of batch per second. Defaults to None. - """ - if name in self.stats_to_track: - self.step = step - self.new_key = not name in self.stats_dict or self.new_key - val = end_time - start_time - if batch_size: - # calculate the batch per second stat - val = batch_size / val - - if step is not None: - # calculate updated average over the buffered max history - curr_history = self.stats_dict.get(name, {"buffer": [], "avg": 0}) - curr_buffer = curr_history["buffer"] - curr_avg = curr_history["avg"] - if len(curr_buffer) >= self.max_history: - curr_buffer.pop(0) - curr_buffer.append(val) - curr_avg = sum(curr_buffer) / len(curr_buffer) - self.stats_dict[name] = {"buffer": curr_buffer, "avg": curr_avg} - else: - # logging total time instead of average - self.stats_dict[name] = val - - 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.graph.max_num_iterations - step - self.stats_dict[Stats.ETA] = remain_iter * self.stats_dict[name]["avg"] - - def handle_header(self): - """helper to handle the printing of the header labels""" - if self.step == 0 or self.new_key: - mssg = f"{'Step (% Done)':<20}" - for k in self.stats_dict: - mssg += f"{k.value:<20} " - if self.step > 0: - cursorup(len(self.past_stats) + 2) - print(mssg) - print("-" * len(mssg)) - if self.step > 0: - for mssg in self.past_stats: - print(mssg) - - def handle_stats(self, fraction_done): - """helper to print out the stats in a readable format""" - if self.step > 0: - curr_mssg = f"{self.step} ({fraction_done*100:.02f}%)" - curr_mssg = f"{curr_mssg:<20}" - for k, v in self.stats_dict.items(): - if isinstance(v, dict): - v = v["avg"] - - if "(time)" in k.value: - v = str(datetime.timedelta(seconds=v)) - elif "(ms)" in k.value: - ms = v * 1e3 - v = f"{ms:0.4f}" - else: - v = f"{v:0.4f}" - curr_mssg += f"{v:<20} " - # update the history buffer - # if len(self.past_stats) >= self.max_history: - # self.past_stats.pop(0) - # cursorup(len(self.past_stats) + 1) - # for mssg in self.past_stats: - # print(mssg) - print(curr_mssg) - # self.past_stats.append(curr_mssg) - - @check_print_stats_step - def print_stats(self, fraction_done: float): - """helper to print out the stats dictionary. - - Args: - fraction_done (float): fraction of steps executed in training iterations - """ - # print a new header line if there is a new key added - # self.handle_header() - - # generate a new stats reporting message - self.handle_stats(fraction_done) - - def dump_stats(self): - """Dump stats locally to a json file""" - raise NotImplementedError - - -def cursorup(x: int): - """utility tool to move the cursor up on the terminal - - Args: - x (int): amount of lines to move cursor upward - """ - print(f"\r\033[{x}A", end="") diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index 735f2687d..3d17de4fb 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -3,9 +3,11 @@ Generic Writer class """ +import datetime +import enum import os from abc import abstractmethod -from typing import Dict +from typing import Any, Dict, List import imageio import numpy as np @@ -15,61 +17,129 @@ from torch.utils.tensorboard import SummaryWriter from torchtyping import TensorType import mattport.utils.writer -from mattport.utils import comms from mattport.utils.decorators import check_main_thread, decorate_all to8b = lambda x: (255 * torch.clamp(x, min=0, max=1)).to(torch.uint8) EVENT_WRITERS = [] +EVENT_STORAGE = [] +GLOBAL_BUFFER = {} -def write_image(name: str, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None) -> None: - """Write image to all writers - Args: - name (str): data identifier - image (TensorType["H", "W", "C"]): rendered image to write - step (int): the time step to log - group (str): the group e.g., "Loss", "Accuracy", "Time" - prefix (str): the prefix e.g., "train-", "test-" - """ - assert comms.is_main_process(), "Writing out with process other than main" - for writer in EVENT_WRITERS: - writer.write_image(name, image, step, group, prefix) - - -def write_scalar(name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None: - """Write scalar to all writers +class EventType(enum.Enum): + """Possible Event types that can be logged Args: - name (str): data identifier - step (int): the time step to log - group (str): the group e.g., "Loss", "Accuracy", "Time" - prefix (str): the prefix e.g., "train-", "test-" + enum (_type_): _description_ """ - assert comms.is_main_process(), "Writing out with process other than main" + + IMAGE = "write_image" + SCALAR = "write_scalar" + DICT = "write_scalar_dict" + + +@check_main_thread +def put_image(name, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None): + EVENT_STORAGE.append( + {"name": name, "write_type": EventType.IMAGE, "event": image, "step": step, "group": group, "prefix": prefix} + ) + + +@check_main_thread +def put_scalar(name: str, scalar: float, step: int, group: str = None, prefix: str = None): + GLOBAL_BUFFER["new_key"] = not name in GLOBAL_BUFFER["events"] or GLOBAL_BUFFER["new_key"] + GLOBAL_BUFFER["events"][name] = scalar + EVENT_STORAGE.append( + {"name": name, "write_type": EventType.SCALAR, "event": scalar, "step": step, "group": group, "prefix": prefix} + ) + + +@check_main_thread +def put_dict(name: str, scalar_dict: float, step: int, group: str = None, prefix: str = None): + EVENT_STORAGE.append( + { + "name": name, + "write_type": EventType.DICT, + "event": scalar_dict, + "step": step, + "group": group, + "prefix": prefix, + } + ) + + +@check_main_thread +def put_time( + name: str, + start_time: float, + end_time: float, + step: int, + group: str = None, + prefix: str = None, + avg_over_iters: bool = False, + avg_over_batch: int = None, + update_eta: bool = False, +): + GLOBAL_BUFFER["step"] = step + GLOBAL_BUFFER["new_key"] = not name in GLOBAL_BUFFER["events"] or GLOBAL_BUFFER["new_key"] + val = end_time - start_time + if avg_over_batch: + val = avg_over_batch / val + + if avg_over_iters: + 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"]: + curr_buffer.pop(0) + curr_buffer.append(val) + curr_avg = sum(curr_buffer) / len(curr_buffer) + put_scalar(name, curr_avg, step, group, prefix) + GLOBAL_BUFFER["events"][name] = {"buffer": curr_buffer, "avg": curr_avg} + else: + GLOBAL_BUFFER["events"][name] = val + put_scalar(name, val, step, group, prefix) + + if update_eta: + ## NOTE: eta should be called with avg train iteration 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) + + +@check_main_thread +def write_out_storage(): for writer in EVENT_WRITERS: - writer.write_scalar(name, scalar, step, group, prefix) - - -def write_scalar_dict(scalar_dict: Dict[str, float], step: int, group: str = None, prefix: str = None) -> None: - """Write dictionary of scalars to all writers - - Args: - scalar_dict (dict): dictionary containing all scalar values with key names and quantities - step (int): the time step to log - group (str): the group e.g., "Loss", "Accuracy", "Time" - prefix (str): the prefix e.g., "train-", "test-" - """ - assert comms.is_main_process(), "Writing out with process other than main" - for writer in EVENT_WRITERS: - writer.write_scalar_dict(scalar_dict, step, group, prefix) + for event in EVENT_STORAGE: + write_func = getattr(writer, event["write_type"].value) + if event == EventType.DICT: + write_func(event["event"], event["step"], event["group"], event["prefix"]) + if isinstance(writer, LocalWriter): + continue + else: + write_func(event["event"], event["name"], event["step"], event["group"], event["prefix"]) + if isinstance(writer, LocalWriter): + break + EVENT_STORAGE.clear() +@check_main_thread def setup_event_writers(config: DictConfig) -> None: """Initialization of all event writers specified in config""" - if comms.is_main_process(): - for writer_type in config: - curr_writer = getattr(mattport.utils.writer, writer_type) - EVENT_WRITERS.append(curr_writer(config[writer_type].save_dir)) + 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] + if writer_type == "LocalWriter": + GLOBAL_BUFFER["max_history"] = curr_config.max_history + GLOBAL_BUFFER["steps_per_log"] = config.logging.steps_per_log + curr_writer = writer_class(curr_config.save_dir, curr_config.stats_to_track, config) + else: + curr_writer = writer_class(curr_config.save_dir) + EVENT_WRITERS.append(curr_writer) + GLOBAL_BUFFER["max_iter"] = config.graph.max_num_iterations + GLOBAL_BUFFER["new_key"] = True + GLOBAL_BUFFER["events"] = {} def get_tensorboard_name(name: str, group: str = None, prefix: str = None): @@ -165,16 +235,94 @@ class TensorboardWriter(Writer): self.tb_writer.add_scalar(tensorboard_name, scalar, step) +def _cursorup(x: int): + """utility tool to move the cursor up on the terminal + + Args: + x (int): amount of lines to move cursor upward + """ + print(f"\r\033[{x}A", end="\x1b[1K\r") + + +def _format_time(v): + if v < 1.0: + ms = v * 1e3 + v = f"{ms:0.4f} ms" + else: + v = str(datetime.timedelta(seconds=v)) + return v + + +def _consolidate_events(): + latest_map = {} + for event in EVENT_STORAGE: + name = event["name"] + latest_map[name] = event["event"] + return latest_map + + @decorate_all([check_main_thread]) class LocalWriter(Writer): """Local Writer Class""" + def __init__(self, save_dir: str, stats_to_track: List[str], config: DictConfig): + super().__init__(save_dir) + self.stats_to_track = stats_to_track + 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 ) -> None: - image = to8b(image) - image_path = os.path.join(self.save_dir, f"{name}.jpg") - imageio.imwrite(image_path, np.uint8(image.cpu().numpy() * 255.0)) + if name in self.stats_to_track: + image = to8b(image) + image_path = os.path.join(self.save_dir, f"{name}.jpg") + imageio.imwrite(image_path, np.uint8(image.cpu().numpy() * 255.0)) def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None: - raise NotImplementedError + latest_map = _consolidate_events() + self.handle_header(latest_map) + self.handle_stats(latest_map) + GLOBAL_BUFFER["new_key"] = False + + 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"]: + 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) + print(mssg) + print("-" * len(mssg)) + 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=" "): + """helper to print out the stats in a readable format""" + step = GLOBAL_BUFFER["step"] + fraction_done = step / GLOBAL_BUFFER["max_iter"] + curr_mssg = f"{step} ({fraction_done*100:.02f}%)" + curr_mssg = f"{curr_mssg:<20}" + for name, v in latest_map.items(): + if name in self.stats_to_track: + if "(time)" in name: + v = _format_time(v) + else: + v = f"{v:0.4f}" + 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) + + for mssg in self.past_stats: + print(f"{mssg:{padding}<{self.max_mssg_len}}") + self.past_stats.append(curr_mssg) From 61238247f28656a37e841426fd440db0d906dc28 Mon Sep 17 00:00:00 2001 From: evonneng Date: Fri, 20 May 2022 15:39:47 +0000 Subject: [PATCH 2/4] basic revamp of stats tracker/writers --- configs/logging/default_logging.yaml | 6 +++--- mattport/nerf/trainer.py | 14 +++++++------- mattport/utils/writer.py | 27 ++++++++++++++++++++------- 3 files changed, 30 insertions(+), 17 deletions(-) diff --git a/configs/logging/default_logging.yaml b/configs/logging/default_logging.yaml index a1279499c..8938d1233 100644 --- a/configs/logging/default_logging.yaml +++ b/configs/logging/default_logging.yaml @@ -1,11 +1,11 @@ steps_per_log: 10 +max_history: 20 # writer logs losses and images per iteration writer: - # TensorboardWriter: - # save_dir: './' + TensorboardWriter: + save_dir: './' LocalWriter: save_dir: './' - max_history: 20 stats_to_track: ["Data Load (time)", "Train Iter (time)", "Rays Per Sec (time)", "Test PSNR"] # profiler logs run times of functions and prints at end of training diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index ee2232ba4..e6f654a24 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -199,7 +199,6 @@ class Trainer: num_iterations = self.config.graph.max_num_iterations iter_dataset = iter(self.train_dataloader) for step in range(self.start_step, self.start_step + num_iterations): - self.step = step data_start = time() batch = next(iter_dataset) writer.put_time( @@ -233,17 +232,18 @@ class Trainer: if step % self.config.graph.steps_per_test == 0: # NOTE(ethan): we should still run this in dry-run mode! for image_idx in self.config.data.val_image_indices: _ = self.test_image(image_idx=image_idx, step=step) - self._write_out_storage() + self._write_out_storage(step) writer.put_time(name="Train Total (time)", start_time=train_start, end_time=time(), step=num_iterations) - self._write_out_storage() + self._write_out_storage(num_iterations) - def _write_out_storage(self): + def _write_out_storage(self, step): """Perform writes only during appropriate time steps""" if ( - 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 + step % self.config.logging.steps_per_log == 0 + or (self.config.graph.steps_per_save and step % self.config.graph.steps_per_save == 0) + or step % self.config.graph.steps_per_test == 0 + or step == self.config.graph.max_num_iterations ): writer.write_out_storage() diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index 3d17de4fb..ecacd7567 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -7,12 +7,12 @@ import datetime import enum import os from abc import abstractmethod -from typing import Any, Dict, List +from typing import Dict import imageio import numpy as np import torch -from omegaconf import DictConfig +from omegaconf import DictConfig, ListConfig from torch.utils.tensorboard import SummaryWriter from torchtyping import TensorType @@ -39,6 +39,7 @@ 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""" EVENT_STORAGE.append( {"name": name, "write_type": EventType.IMAGE, "event": image, "step": step, "group": group, "prefix": prefix} ) @@ -46,6 +47,7 @@ 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""" GLOBAL_BUFFER["new_key"] = not name in GLOBAL_BUFFER["events"] or GLOBAL_BUFFER["new_key"] GLOBAL_BUFFER["events"][name] = scalar EVENT_STORAGE.append( @@ -55,6 +57,7 @@ def put_scalar(name: str, scalar: float, step: int, group: str = None, prefix: s @check_main_thread def put_dict(name: str, scalar_dict: float, step: int, group: str = None, prefix: str = None): + """Setter function to place a dictionary of scalars into the queue to be written out""" EVENT_STORAGE.append( { "name": name, @@ -79,6 +82,13 @@ def put_time( avg_over_batch: int = None, update_eta: bool = False, ): + """Setter function to place a time element into the queue to be written out + + Processes the time info according to the options: + avg_over_iters (bool): if True, calculate and record a running average of the times + 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 + """ GLOBAL_BUFFER["step"] = step GLOBAL_BUFFER["new_key"] = not name in GLOBAL_BUFFER["events"] or GLOBAL_BUFFER["new_key"] val = end_time - start_time @@ -109,15 +119,16 @@ def put_time( @check_main_thread def write_out_storage(): + """Function that writes all the events in storage to all the writer locations""" for writer in EVENT_WRITERS: for event in EVENT_STORAGE: write_func = getattr(writer, event["write_type"].value) - if event == EventType.DICT: + if event["write_type"] == EventType.DICT: write_func(event["event"], event["step"], event["group"], event["prefix"]) if isinstance(writer, LocalWriter): continue else: - write_func(event["event"], event["name"], event["step"], event["group"], event["prefix"]) + write_func(event["name"], event["event"], event["step"], event["group"], event["prefix"]) if isinstance(writer, LocalWriter): break EVENT_STORAGE.clear() @@ -131,13 +142,15 @@ def setup_event_writers(config: DictConfig) -> None: writer_class = getattr(mattport.utils.writer, writer_type) curr_config = logging_configs[writer_type] if writer_type == "LocalWriter": - GLOBAL_BUFFER["max_history"] = curr_config.max_history - GLOBAL_BUFFER["steps_per_log"] = config.logging.steps_per_log curr_writer = writer_class(curr_config.save_dir, curr_config.stats_to_track, config) else: curr_writer = writer_class(curr_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["steps_per_log"] = config.logging.steps_per_log GLOBAL_BUFFER["new_key"] = True GLOBAL_BUFFER["events"] = {} @@ -265,7 +278,7 @@ def _consolidate_events(): class LocalWriter(Writer): """Local Writer Class""" - def __init__(self, save_dir: str, stats_to_track: List[str], config: DictConfig): + def __init__(self, save_dir: str, stats_to_track: ListConfig, config: DictConfig): super().__init__(save_dir) self.stats_to_track = stats_to_track self.past_stats = [] From bbfe815504baf9b46c9f3b74c5c8e31b7fc22048 Mon Sep 17 00:00:00 2001 From: evonneng Date: Fri, 20 May 2022 18:38:18 +0000 Subject: [PATCH 3/4] 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) From 2f4fa3ff9fd38ad311c7039376bef08893d81107 Mon Sep 17 00:00:00 2001 From: evonneng Date: Fri, 20 May 2022 18:59:00 +0000 Subject: [PATCH 4/4] update README logging --- README.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/README.md b/README.md index 20ffddb2b..e0bc7e199 100644 --- a/README.md +++ b/README.md @@ -56,9 +56,7 @@ python scripts/run_train.py data/dataset=friends_TBBT-big_living_room graph.netw 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. - -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. +1. **Writer**: Logs losses and generated images during training to a specified output stream. Specify the type of writer (Tensorboard, Local Writer, Weights and Biases), and how often to log 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.