From 43d588f35f8d0746953bc47cc16c0cc234a814df Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Wed, 28 Sep 2022 10:10:18 -0700 Subject: [PATCH] specify a viewers list (#574) * viewer list * before master merge * update script --- nerfactory/configs/base.py | 21 +++++++++++++++------ nerfactory/configs/base_configs.py | 11 ++++------- nerfactory/engine/trainer.py | 20 ++++++++++++-------- nerfactory/utils/decorators.py | 4 ++-- nerfactory/utils/writer.py | 24 ++++++++++++++++-------- nerfactory/viewer/server/viewer_utils.py | 8 +++----- scripts/viewer/view_dataset.py | 4 ++-- tests/test_run_train.py | 2 +- 8 files changed, 55 insertions(+), 39 deletions(-) diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 27ab47cdb..6f32faa4b 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -17,7 +17,7 @@ from __future__ import annotations import logging -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Literal, Optional, Tuple, Type @@ -118,8 +118,6 @@ class LoggingConfig(PrintableConfig): max_buffer_size: int = 20 """maximum history size to keep for computing running averages of stats. e.g. if 20, averages will be computed over past 20 occurances.""" - event_writer: Literal["tb", "wandb", "none"] = "wandb" - """specify which writer to use (tensorboard or wandb)""" local_writer: LocalWriterConfig = LocalWriterConfig(enable=True) """if provided, will print stats locally. if None, will disable printing""" enable_profiler: bool = True @@ -154,7 +152,6 @@ class TrainerConfig(PrintableConfig): load_step: Optional[int] = None """optionally specify model step to load from; if none, will find most recent model in load_dir""" load_config: Optional[Path] = None - """optionally specify a pre-defined config to load from""" # Viewer related configs @@ -162,8 +159,6 @@ class TrainerConfig(PrintableConfig): class ViewerConfig(PrintableConfig): """Configuration for viewer instantiation""" - enable: bool = False - """whether to enable viewer""" relative_log_filename: str = "viewer_log_filename.txt" """Filename to use for the log file.""" start_train: bool = True @@ -255,6 +250,20 @@ class Config(PrintableConfig): } } ) + """optionally specify a pre-defined config to load from""" + vis: List[Literal["viewer", "wandb", "tensorboard"]] = field(default_factory=lambda: ["wandb"]) + + def is_viewer_enabled(self) -> bool: + """Checks if a viewer is enabled.""" + return "viewer" in self.vis + + def is_wandb_enabled(self) -> bool: + """Checks if wandb is enabled.""" + return "wandb" in self.vis + + def is_tensorboard_enabled(self) -> bool: + """Checks if tensorboard is enabled.""" + return "tensorboard" in self.vis def set_timestamp(self) -> None: """Dynamically set the experiment timestamp""" diff --git a/nerfactory/configs/base_configs.py b/nerfactory/configs/base_configs.py index 60aea1a13..a5a5ba9c4 100644 --- a/nerfactory/configs/base_configs.py +++ b/nerfactory/configs/base_configs.py @@ -25,7 +25,6 @@ import dcargs from nerfactory.configs.base import ( AdamOptimizerConfig, Config, - LoggingConfig, RAdamOptimizerConfig, SchedulerConfig, TrainerConfig, @@ -64,8 +63,7 @@ base_configs["compound"] = Config( "scheduler": None, } }, - viewer=ViewerConfig(enable=True), - logging=LoggingConfig(event_writer="none"), + vis=["viewer"], ) base_configs["instant-ngp"] = Config( @@ -81,8 +79,7 @@ base_configs["instant-ngp"] = Config( "scheduler": None, } }, - viewer=ViewerConfig(enable=True), - logging=LoggingConfig(event_writer="none"), + vis=["viewer"], ) base_configs["proposal"] = Config( @@ -104,8 +101,8 @@ base_configs["proposal"] = Config( "scheduler": None, }, }, - viewer=ViewerConfig(enable=True, num_rays_per_chunk=2 << 15), - logging=LoggingConfig(event_writer="none"), + viewer=ViewerConfig(num_rays_per_chunk=2 << 15), + vis=["viewer"], ) base_configs["mipnerf-360"] = Config( diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index e133ac98d..9f207bc11 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -106,11 +106,14 @@ class Trainer: self.prev_ckpt_paths = [] # set up viewer if enabled viewer_log_path = self.base_dir / config.viewer.relative_log_filename - self.viewer_state, banner_messages = viewer_utils.setup_viewer(config.viewer, log_filename=viewer_log_path) + if self.config.is_viewer_enabled(): + self.viewer_state, banner_messages = viewer_utils.setup_viewer(config.viewer, log_filename=viewer_log_path) + else: + self.viewer_state, banner_messages = None, None self._check_viewer_warnings() # set up writers/profilers if enabled writer_log_path = self.base_dir / config.logging.relative_log_dir - writer.setup_event_writer(config.logging, log_dir=writer_log_path) + writer.setup_event_writer(config, log_dir=writer_log_path) writer.setup_local_writer( config.logging, max_iter=config.trainer.max_num_iterations, banner_messages=banner_messages ) @@ -188,17 +191,18 @@ class Trainer: def _check_viewer_warnings(self) -> None: """Helper to print out any warnings regarding the way the viewer/loggers are enabled""" - if self.config.viewer.enable: - if self.config.logging.event_writer == "none": + if self.config.is_viewer_enabled(): + is_logger_enabled = self.config.is_wandb_enabled() or self.config.is_tensorboard_enabled() + if is_logger_enabled: string = ( - "[WARNING] Disabling eval iterations since viewer is enabled." - "Please set `--logging.event_writer wandb` (or tb) to run evaluations." + "[WARNING]: Tensorboard or Wandb enabled with Viewer will slow down Viewer. " + "Please set `--vis viewer` for faster rendering." ) CONSOLE.print(f"[bold red]{string}") else: string = ( - "[WARNING]: Tensorboard or Wandb enabled with Viewer will slow down Viewer. " - "Please set `--logging.event_writer none` for faster rendering" + "[WARNING] Not running eval iterations since only viewer is enabled." + " Please add 'wandb' or 'tensorboard' to the `--vis` list to run evaluations." ) CONSOLE.print(f"[bold red]{string}") diff --git a/nerfactory/utils/decorators.py b/nerfactory/utils/decorators.py index cb28dca66..b53551d57 100644 --- a/nerfactory/utils/decorators.py +++ b/nerfactory/utils/decorators.py @@ -54,7 +54,7 @@ def check_viewer_enabled(func: Callable) -> Callable: def wrapper(self, *args, **kwargs): ret = None - if self.config.viewer.enable: + if self.config.is_viewer_enabled(): ret = func(self, *args, **kwargs) return ret @@ -66,7 +66,7 @@ def check_eval_enabled(func: Callable) -> Callable: def wrapper(self, *args, **kwargs): ret = None - if self.config.logging.event_writer != "none": + if self.config.is_wandb_enabled() or self.config.is_tensorboard_enabled(): ret = func(self, *args, **kwargs) return ret diff --git a/nerfactory/utils/writer.py b/nerfactory/utils/writer.py index 999b38597..e65c7803f 100644 --- a/nerfactory/utils/writer.py +++ b/nerfactory/utils/writer.py @@ -26,6 +26,7 @@ from typing import Any, Dict, List, Optional, Union import torch import wandb +from rich import console from torch.utils.tensorboard import SummaryWriter from torchtyping import TensorType @@ -33,6 +34,8 @@ from nerfactory.configs import base as cfg from nerfactory.utils.decorators import check_main_thread, decorate_all from nerfactory.utils.misc import human_format +CONSOLE = console.Console() + to8b = lambda x: (255 * torch.clamp(x, min=0, max=1)).to(torch.uint8) EVENT_WRITERS = [] EVENT_STORAGE = [] @@ -187,7 +190,7 @@ def setup_local_writer(config: cfg.LoggingConfig, max_iter: int, banner_messages GLOBAL_BUFFER["events"] = {} -def setup_event_writer(config: cfg.LoggingConfig, log_dir: Path) -> None: +def setup_event_writer(config: cfg.Config, log_dir: Path) -> None: """Initialization of all event writers specified in config Args: @@ -195,15 +198,20 @@ def setup_event_writer(config: cfg.LoggingConfig, log_dir: Path) -> None: max_iter: maximum number of train iterations banner_messages: list of messages to always display at bottom of screen """ - if config.event_writer == "tb": - curr_writer = TensorboardWriter(log_dir=log_dir) - elif config.event_writer == "wandb": + using_event_writer = False + if config.is_wandb_enabled(): curr_writer = WandbWriter(log_dir=log_dir) + EVENT_WRITERS.append(curr_writer) + using_event_writer = True + if config.is_tensorboard_enabled(): + curr_writer = TensorboardWriter(log_dir=log_dir) + EVENT_WRITERS.append(curr_writer) + using_event_writer = True + if using_event_writer: + string = f"logging events to: {log_dir}" else: - logging.info("disabled tensorboard/wandb event writers") - return - EVENT_WRITERS.append(curr_writer) - logging.info("logging %s events to: %s", config.event_writer, log_dir) + string = "disabled tensorboard/wandb event writers" + CONSOLE.print(f"[bold red]{string}") class Writer: diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index 03ebe8099..8bdbc7b09 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/nerfactory/viewer/server/viewer_utils.py @@ -50,11 +50,9 @@ def setup_viewer(config: cfg.ViewerConfig, log_filename: Path): Args: config: the configuration to instantiate viewer """ - if config.enable: - viewer_state = ViewerState(config, log_filename=log_filename) - banner_messages = [f"Viewer at: {viewer_state.viewer_url}"] - return viewer_state, banner_messages - return None, None + viewer_state = ViewerState(config, log_filename=log_filename) + banner_messages = [f"Viewer at: {viewer_state.viewer_url}"] + return viewer_state, banner_messages class OutputTypes(str, enum.Enum): diff --git a/scripts/viewer/view_dataset.py b/scripts/viewer/view_dataset.py index 3f894a2d9..682db3b22 100755 --- a/scripts/viewer/view_dataset.py +++ b/scripts/viewer/view_dataset.py @@ -26,8 +26,8 @@ torch.backends.cudnn.benchmark = True # type: ignore def main(config: cfg.Config) -> None: """Main function.""" - if not config.viewer.enable: - config.viewer.enable = True + if not config.is_viewer_enabled(): + config.vis.append("viewer") logging.info("Enabling viewer to view dataset") viewer_state = viewer_utils.ViewerState(config.viewer) datamanager = config.pipeline.datamanager.setup() diff --git a/tests/test_run_train.py b/tests/test_run_train.py index e6f97db5b..e15c0aa48 100644 --- a/tests/test_run_train.py +++ b/tests/test_run_train.py @@ -27,7 +27,7 @@ def set_reduced_config(config: Config): config.pipeline.datamanager.train_num_rays_per_batch = 4 # use tensorboard logging instead of wandb - config.logging.event_writer = "tb" + config.vis = ["tensorboard"] config.logging.relative_log_dir = Path("/tmp/") # reduce model factors