From 0c43e7c4b27f68ceed13de4e0dbbf226663612f7 Mon Sep 17 00:00:00 2001 From: evonneng Date: Wed, 21 Sep 2022 13:40:56 -0700 Subject: [PATCH] updated event writing logic/moved save and print to after config set (#502) * (1) if wandb/tb enabled always eval; (2) move printing of config to after all variables set * (3) save config after everything is set * dont change config; just disable eval with decorator --- nerfactory/configs/base.py | 19 ++++++++ nerfactory/engine/trainer.py | 84 ++++++++++++++++++++-------------- nerfactory/utils/decorators.py | 12 +++++ scripts/run_train.py | 15 ++---- 4 files changed, 83 insertions(+), 47 deletions(-) diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 5e0263ef8..7b073deab 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -16,6 +16,7 @@ from __future__ import annotations +import logging from dataclasses import dataclass from datetime import datetime from pathlib import Path @@ -23,6 +24,8 @@ from typing import Any, Dict, List, Literal, Optional, Tuple, Type import dcargs import torch +import yaml +from rich.console import Console from nerfactory.configs.utils import to_immutable_dict @@ -411,3 +414,19 @@ class Config(PrintableConfig): assert self.method_name is not None, "Please set method name in config or via the cli" self.set_experiment_name() return Path(f"outputs/{self.experiment_name}/{self.method_name}/{self.timestamp}") + + def print_to_terminal(self) -> None: + """Helper to pretty print config to terminal""" + console = Console(width=120) + console.rule("Config") + console.print(self) + console.rule("") + + def save_config(self) -> None: + """Save config to base directory""" + base_dir = self.get_base_dir() + assert base_dir is not None + base_dir.mkdir(parents=True, exist_ok=True) + config_yaml_path = base_dir / "config.yml" + logging.info(f"Saving config to: {config_yaml_path}") + config_yaml_path.write_text(yaml.dump(self), "utf8") diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index a4f9329d9..3f3ba7923 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -38,7 +38,11 @@ from nerfactory.utils.callbacks import ( TrainingCallbackAttributes, TrainingCallbackLocation, ) -from nerfactory.utils.decorators import check_main_thread, check_viewer_enabled +from nerfactory.utils.decorators import ( + check_eval_enabled, + check_main_thread, + check_viewer_enabled, +) from nerfactory.utils.misc import step_check from nerfactory.utils.writer import EventName, TimeWriter from nerfactory.viewer.server import viewer_utils @@ -143,7 +147,6 @@ class Trainer: with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME): num_iterations = self.config.trainer.max_num_iterations for step in range(self._start_step, self._start_step + num_iterations): - # if the viewer used, the rendering of the viewer will be included in the iteration train time with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as train_t: self.pipeline.train() @@ -176,33 +179,7 @@ class Trainer: writer.put_dict(name="Train Loss Dict", scalar_dict=loss_dict, step=step) writer.put_dict(name="Train Metrics Dict", scalar_dict=metrics_dict, step=step) - # a batch of eval rays - if step_check(step, self.config.trainer.steps_per_eval_batch, run_at_zero=True): - _, eval_loss_dict, eval_metrics_dict = self.pipeline.get_eval_loss_dict(step=step) - eval_loss = functools.reduce(torch.add, eval_loss_dict.values()) - writer.put_scalar(name="Eval Loss", scalar=eval_loss, step=step) - writer.put_dict(name="Eval Loss Dict", scalar_dict=eval_loss_dict, step=step) - writer.put_dict(name="Eval Metrics Dict", scalar_dict=eval_metrics_dict, step=step) - - # one eval image - if step_check(step, self.config.trainer.steps_per_eval_image): - with TimeWriter(writer, EventName.TEST_RAYS_PER_SEC, write=False) as test_t: - metrics_dict, images_dict = self.pipeline.get_eval_image_metrics_and_images(step=step) - writer.put_time( - name=EventName.TEST_RAYS_PER_SEC, - duration=metrics_dict["num_rays"] / test_t.duration, - step=step, - avg_over_steps=True, - ) - writer.put_dict(name="Eval Images Metrics", scalar_dict=metrics_dict, step=step) - group = "Eval Images" - for image_name, image in images_dict.items(): - writer.put_image(name=group + "/" + image_name, image=image, step=step) - - # all eval images - if step_check(step, self.config.trainer.steps_per_eval_all_images): - metrics_dict = self.pipeline.get_average_eval_image_metrics(step=step) - writer.put_dict(name="Eval Images Metrics Dict (all images)", scalar_dict=metrics_dict, step=step) + self.eval_iteration(step) if step != 0 and self.config.trainer.steps_per_save and step % self.config.trainer.steps_per_save == 0: self._save_checkpoint(step) @@ -212,17 +189,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.logging.event_writer == "none": + string = ( + "[WARNING] Disabling eval iterations since viewer is enabled." + "Please set `--logging.event_writer wandb` (or tb) to run evaluations." + ) + 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" ) CONSOLE.print(f"[bold red]{string}") - self.config.trainer.steps_per_eval_batch = self.config.trainer.max_num_iterations - self.config.trainer.steps_per_eval_image = self.config.trainer.max_num_iterations - self.config.trainer.steps_per_eval_all_images = self.config.trainer.max_num_iterations - string = "[WARNING] Disabling eval iterations since viewer is enabled." - CONSOLE.print(f"[bold red]{string}") @check_viewer_enabled def _init_viewer_scene(self) -> None: @@ -334,3 +312,39 @@ class Trainer: # Merging loss and metrics dict into a single output. return loss, loss_dict, metrics_dict + + @check_eval_enabled + @profiler.time_function + def eval_iteration(self, step): + """Run one iteration with different batch/image/all image evaluations depending on step size. + + Args: + step: Current training step. + """ + # a batch of eval rays + if step_check(step, self.config.trainer.steps_per_eval_batch, run_at_zero=True): + _, eval_loss_dict, eval_metrics_dict = self.pipeline.get_eval_loss_dict(step=step) + eval_loss = functools.reduce(torch.add, eval_loss_dict.values()) + writer.put_scalar(name="Eval Loss", scalar=eval_loss, step=step) + writer.put_dict(name="Eval Loss Dict", scalar_dict=eval_loss_dict, step=step) + writer.put_dict(name="Eval Metrics Dict", scalar_dict=eval_metrics_dict, step=step) + + # one eval image + if step_check(step, self.config.trainer.steps_per_eval_image): + with TimeWriter(writer, EventName.TEST_RAYS_PER_SEC, write=False) as test_t: + metrics_dict, images_dict = self.pipeline.get_eval_image_metrics_and_images(step=step) + writer.put_time( + name=EventName.TEST_RAYS_PER_SEC, + duration=metrics_dict["num_rays"] / test_t.duration, + step=step, + avg_over_steps=True, + ) + writer.put_dict(name="Eval Images Metrics", scalar_dict=metrics_dict, step=step) + group = "Eval Images" + for image_name, image in images_dict.items(): + writer.put_image(name=group + "/" + image_name, image=image, step=step) + + # all eval images + if step_check(step, self.config.trainer.steps_per_eval_all_images): + metrics_dict = self.pipeline.get_average_eval_image_metrics(step=step) + writer.put_dict(name="Eval Images Metrics Dict (all images)", scalar_dict=metrics_dict, step=step) diff --git a/nerfactory/utils/decorators.py b/nerfactory/utils/decorators.py index 77524f699..cb28dca66 100644 --- a/nerfactory/utils/decorators.py +++ b/nerfactory/utils/decorators.py @@ -61,6 +61,18 @@ def check_viewer_enabled(func: Callable) -> Callable: return wrapper +def check_eval_enabled(func: Callable) -> Callable: + """Decorator: check if evaluation step is enabled""" + + def wrapper(self, *args, **kwargs): + ret = None + if self.config.logging.event_writer != "none": + ret = func(self, *args, **kwargs) + return ret + + return wrapper + + def check_main_thread(func: Callable) -> Callable: """Decorator: check if you are on main thread""" diff --git a/scripts/run_train.py b/scripts/run_train.py index a373b9f69..472b0b67c 100755 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -16,7 +16,6 @@ import torch import torch.distributed as dist import torch.multiprocessing as mp import yaml -from rich.console import Console from nerfactory.configs import base as cfg from nerfactory.configs.base_configs import AnnotatedBaseConfigUnion @@ -189,17 +188,9 @@ def main(config: cfg.Config) -> None: logging.info(f"Loading pre-set config from: {config.trainer.load_config}") config = yaml.load(config.trainer.load_config.read_text(), Loader=yaml.Loader) - console = Console(width=120) - console.rule("Config") - console.print(config) - console.rule("") - - base_dir = config.get_base_dir() - assert base_dir is not None - base_dir.mkdir(parents=True, exist_ok=True) - config_yaml_path = base_dir / "config.yml" - logging.info(f"Saving config to: {config_yaml_path}") - config_yaml_path.write_text(yaml.dump(config), "utf8") + # print and save config + config.print_to_terminal() + config.save_config() launch( train_loop,