From 61238247f28656a37e841426fd440db0d906dc28 Mon Sep 17 00:00:00 2001 From: evonneng Date: Fri, 20 May 2022 15:39:47 +0000 Subject: [PATCH] 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 = []