diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 6daef9911..23e16ad35 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -17,7 +17,7 @@ from dataclasses import dataclass from datetime import datetime from pathlib import Path -from typing import Any, Dict, Literal, Optional, Tuple, Type +from typing import Any, Dict, List, Literal, Optional, Tuple, Type import torch @@ -110,6 +110,14 @@ class LocalWriterConfig(InstantiateConfig): relative_log_dir: Path = Path("./") log_dir: Optional[Path] = None # full log dir path to be dynamically set + def setup(self, banner_messages: Optional[List[str]] = None, **kwargs) -> Any: + """Instantiate local writer + + Args: + banner_messages: List of strings that always print at the bottom of screen. Defaults to None. + """ + return self._target(self, banner_messages=banner_messages, **kwargs) + @dataclass class LoggingConfig(PrintableConfig): diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index fe79cf13c..4d6de5a5c 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -83,14 +83,19 @@ class Trainer: self.pipeline: Pipeline self.optimizers: Optimizers self.start_step = 0 - # logging variables - writer.setup_event_writers(config.logging, max_iter=config.trainer.max_num_iterations) - profiler.setup_profiler(config.logging) # visualizer variable + banner_messages = None self.visualizer_state = viewer_utils.VisualizerState(config.viewer) + if config.viewer.enable: + banner_messages = [f"Viewer at: {self.visualizer_state.viewer_url}"] self.grad_scaler = GradScaler(enabled=self.mixed_precision) # training callbacks self.callbacks: List[Callback] + # logging variables + writer.setup_event_writers( + config.logging, max_iter=config.trainer.max_num_iterations, banner_messages=banner_messages + ) + profiler.setup_profiler(config.logging) def setup(self, test_mode=False): """Setup the Trainer by calling other setup functions. diff --git a/nerfactory/utils/writer.py b/nerfactory/utils/writer.py index 64ee430f8..bbb898d2e 100644 --- a/nerfactory/utils/writer.py +++ b/nerfactory/utils/writer.py @@ -22,7 +22,7 @@ import logging from abc import abstractmethod from pathlib import Path from time import time -from typing import Dict, Optional +from typing import Dict, List, Optional import imageio import numpy as np @@ -137,11 +137,14 @@ def write_out_storage(): @check_main_thread -def setup_event_writers(config: cfg.LoggingConfig, max_iter: int) -> None: +def setup_event_writers(config: cfg.LoggingConfig, max_iter: int, banner_messages: Optional[List[str]] = None) -> None: """Initialization of all event writers specified in config""" for writer_type_config in config.writer: if writer_type_config.enable: - curr_writer = writer_type_config.setup() + if isinstance(writer_type_config, cfg.LocalWriterConfig): + curr_writer = writer_type_config.setup(banner_messages=banner_messages) + else: + curr_writer = writer_type_config.setup() EVENT_WRITERS.append(curr_writer) logging.info("logging info to: %s", writer_type_config.log_dir) @@ -309,7 +312,7 @@ def _format_time(seconds): class LocalWriter(Writer): """Local Writer Class""" - def __init__(self, config: cfg.LocalWriterConfig): + def __init__(self, config: cfg.LocalWriterConfig, banner_messages: Optional[List[str]] = None): """ Args: stats_to_track (ListConfig): the names of stats that should be logged. @@ -320,6 +323,9 @@ class LocalWriter(Writer): self.max_log_size = config.max_log_size self.keys = set() self.past_mssgs = ["", ""] + self.banner_len = 0 if banner_messages is None else len(banner_messages) + if banner_messages: + self.past_mssgs.extend(banner_messages) self.has_printed = False def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: @@ -330,6 +336,11 @@ class LocalWriter(Writer): def write_scalar(self, name: str, scalar: float, step: int) -> None: if step > 0: + if not self.has_printed and self.max_log_size: + logging.info( + "Printing max of %d lines. Set flag `--logging.writer.2.max-log-size=0` to disable line wrapping.", + self.max_log_size, + ) latest_map, new_key = self._consolidate_events() self._update_header(latest_map, new_key) self._print_stats(latest_map) @@ -348,7 +359,7 @@ class LocalWriter(Writer): def _update_header(self, latest_map, new_key): """helper to handle the printing of the header labels""" full_log_cond = not self.max_log_size and GLOBAL_BUFFER["step"] <= GLOBAL_BUFFER["steps_per_log"] - capped_log_cond = self.max_log_size and (len(self.past_mssgs) <= 2 or new_key) + capped_log_cond = self.max_log_size and (len(self.past_mssgs) - self.banner_len <= 2 or new_key) if full_log_cond or capped_log_cond: mssg = f"{'Step (% Done)':<20}" for name, _ in latest_map.items(): @@ -359,7 +370,7 @@ class LocalWriter(Writer): if full_log_cond or not self.has_printed: print(mssg) print("-" * len(mssg)) - self.has_printed = True + # self.has_printed = True def _print_stats(self, latest_map, padding=" "): """helper to print out the stats in a readable format""" @@ -379,14 +390,19 @@ class LocalWriter(Writer): # update the history buffer if self.max_log_size: - cursor_idx = len(self.past_mssgs) - if len(self.past_mssgs[2:]) >= self.max_log_size: + if not self.has_printed: + cursor_idx = len(self.past_mssgs) - self.banner_len + self.has_printed = True + else: + cursor_idx = len(self.past_mssgs) + if len(self.past_mssgs[2:]) - self.banner_len >= self.max_log_size: self.past_mssgs.pop(2) - self.past_mssgs.append(curr_mssg) + self.past_mssgs.insert(len(self.past_mssgs) - self.banner_len, curr_mssg) _cursorup(cursor_idx) - for mssg in self.past_mssgs: - pad_len = len(self.past_mssgs[0]) - print(f"{mssg:{padding}<{pad_len}}") + for i, mssg in enumerate(self.past_mssgs): + pad_len = len(max(self.past_mssgs, key=len)) + style = "\033[2;30;45m" if self.banner_len and i >= len(self.past_mssgs) - self.banner_len else "" + print(f"{style}{mssg:{padding}<{pad_len}} \033[0;0m") else: print(curr_mssg) diff --git a/scripts/run_train.py b/scripts/run_train.py index 7bd071a4e..f474ba65c 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -208,5 +208,4 @@ if __name__ == "__main__": config_yaml_path = instantiated_config.base_dir / "config.yml" logging.info(f"Saving config to: {config_yaml_path}") config_yaml_path.write_text(yaml.dump(instantiated_config), "utf8") - main(config=instantiated_config)