diff --git a/configs/logging/default_logging.yaml b/configs/logging/default_logging.yaml index 59ef1333b..8938d1233 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 +max_history: 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"] + # 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 5eb4f7fe7..f1ba247ba 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 @@ -232,9 +232,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, ...] @@ -246,13 +246,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..e6f654a24 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): 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(step) - 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(num_iterations) + + def _write_out_storage(self, step): + """Perform writes only during appropriate time steps""" + if ( + 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() @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..ecacd7567 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -3,6 +3,8 @@ Generic Writer class """ +import datetime +import enum import os from abc import abstractmethod from typing import Dict @@ -10,66 +12,147 @@ 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 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" - for writer in EVENT_WRITERS: - writer.write_scalar(name, scalar, step, group, prefix) + + IMAGE = "write_image" + SCALAR = "write_scalar" + DICT = "write_scalar_dict" -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 +@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} + ) - 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-" + +@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( + {"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): + """Setter function to place a dictionary of scalars into the queue to be written out""" + 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, +): + """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 """ - assert comms.is_main_process(), "Writing out with process other than main" + 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(): + """Function that writes all the events in storage to all the writer locations""" 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["write_type"] == EventType.DICT: + write_func(event["event"], event["step"], event["group"], event["prefix"]) + if isinstance(writer, LocalWriter): + continue + else: + write_func(event["name"], event["event"], 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": + 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"] = {} def get_tensorboard_name(name: str, group: str = None, prefix: str = None): @@ -165,16 +248,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: ListConfig, 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)