From 0506de65cd6c8aa3f9ab8e3af9d328f7c88015f1 Mon Sep 17 00:00:00 2001 From: evonneng Date: Mon, 10 Oct 2022 09:42:15 -0700 Subject: [PATCH] fixed --start-train flag (#741) * wip; fixing --start-train flag * fixed no rendering --- nerfstudio/engine/trainer.py | 13 +++++++++---- nerfstudio/viewer/server/viewer_utils.py | 12 +++++++++--- 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index ddb8772d7..4a6546de7 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -190,10 +190,13 @@ class Trainer: writer.write_out_storage() # save checkpoint at the end of training self.save_checkpoint(step) - if self.config.is_viewer_enabled(): - while True: - self.viewer_state.vis["renderingState/isTraining"].write(False) - self._update_viewer_state(step) + self._always_render(step) + + def _always_render(self, step): + if self.config.is_viewer_enabled(): + while True: + self.viewer_state.vis["renderingState/isTraining"].write(False) + self._update_viewer_state(step) def _check_viewer_warnings(self) -> None: """Helper to print out any warnings regarding the way the viewer/loggers are enabled""" @@ -212,6 +215,8 @@ class Trainer: dataset=self.pipeline.datamanager.train_dataset, start_train=self.config.viewer.start_train, ) + if not self.config.viewer.start_train: + self._always_render(self._start_step) @check_viewer_enabled def _update_viewer_state(self, step: int): diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 41d2fd1fb..893845a3e 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -485,7 +485,10 @@ class ViewerState: # process remaining training ETA self.vis["renderingState/train_eta"].write(GLOBAL_BUFFER["events"].get(EventName.ETA.value, "Starting")) # process ratio time spent on vis vs train - if EventName.ITER_VIS_TIME.value in GLOBAL_BUFFER["events"]: + if ( + EventName.ITER_VIS_TIME.value in GLOBAL_BUFFER["events"] + and EventName.ITER_TRAIN_TIME.value in GLOBAL_BUFFER["events"] + ): vis_time = GLOBAL_BUFFER["events"][EventName.ITER_VIS_TIME.value]["avg"] train_time = GLOBAL_BUFFER["events"][EventName.ITER_TRAIN_TIME.value]["avg"] vis_train_ratio = f"{int(vis_time / train_time * 100)}% spent on viewer" @@ -519,8 +522,12 @@ class ViewerState: if EventName.TRAIN_RAYS_PER_SEC.value in GLOBAL_BUFFER["events"]: train_rays_per_sec = GLOBAL_BUFFER["events"][EventName.TRAIN_RAYS_PER_SEC.value]["avg"] + elif not is_training: + train_rays_per_sec = ( + 80000 # TODO(eventually find a way to not hardcode. case where there are no prior training steps) + ) else: - return None + return None, None if EventName.VIS_RAYS_PER_SEC.value in GLOBAL_BUFFER["events"]: vis_rays_per_sec = GLOBAL_BUFFER["events"][EventName.VIS_RAYS_PER_SEC.value]["avg"] else: @@ -540,7 +547,6 @@ class ViewerState: if image_width > self.max_resolution: image_width = self.max_resolution image_height = int(image_width / aspect_ratio) - return image_height, image_width def _process_invalid_output(self, output_type: str) -> str: