diff --git a/docs/tooling/logging_profiling.md b/docs/tooling/logging_profiling.md index 6b6ff77e8..d55ca7607 100644 --- a/docs/tooling/logging_profiling.md +++ b/docs/tooling/logging_profiling.md @@ -18,7 +18,7 @@ logging: LocalWriter: log_dir: "./" stats_to_track: - [ITER_TRAIN_TIME, RAYS_PER_SEC, CURR_TEST_PSNR] + [ITER_TRAIN_TIME, TRAIN_RAYS_PER_SEC, CURR_TEST_PSNR] max_log_size: 10 ``` @@ -46,7 +46,7 @@ WandbWriter: LocalWriter: log_dir: "./" # if images are saved, where the images would be saved stats_to_track: - [ITER_TRAIN_TIME, RAYS_PER_SEC, CURR_TEST_PSNR] + [ITER_TRAIN_TIME, TRAIN_RAYS_PER_SEC, CURR_TEST_PSNR] max_log_size: 10 ``` diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 65674b957..17f6eaf49 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -115,8 +115,10 @@ class LocalWriterConfig(InstantiateConfig): enable: bool = False stats_to_track: Tuple[writer.EventName, ...] = ( writer.EventName.ITER_TRAIN_TIME, - writer.EventName.RAYS_PER_SEC, + writer.EventName.TRAIN_RAYS_PER_SEC, writer.EventName.CURR_TEST_PSNR, + writer.EventName.VIS_RAYS_PER_SEC, + writer.EventName.TEST_RAYS_PER_SEC, ) max_log_size: int = 10 relative_log_dir: Path = Path("./") diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index 822948b06..4822507b1 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -122,7 +122,7 @@ class Trainer: for step in range(self.start_step, self.start_step + num_iterations): # if the visualizer used, the rendering of the visualizer will be included in the iteration train time - with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as t: + with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as train_t: # training callbacks before the training iteration for callback in self.callbacks: @@ -136,22 +136,36 @@ class Trainer: for callback in self.callbacks: callback.run_callback_at_location(step, location=TrainingCallbackLocation.AFTER_TRAIN_ITERATION) - with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as _: + with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as vis_t: self.visualizer_state.update_scene(step, self.pipeline.model) - train_num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch - writer.put_scalar(name=EventName.RAYS_PER_SEC, scalar=train_num_rays_per_batch / t.duration, step=step) - if step != 0 and step % self.config.logging.steps_per_log == 0: writer.put_dict(name="Loss/train-loss_metrics_dict", scalar_dict=loss_metric_dict, step=step) if step != 0 and self.config.trainer.steps_per_save and step % self.config.trainer.steps_per_save == 0: self._save_checkpoint(self.config.trainer.model_dir, step) if step % self.config.trainer.steps_per_test == 0: self.pipeline.get_eval_loss_dict(step=step) + self._update_rays_per_sec(train_t, vis_t, step) self._write_out_storage(step) self._write_out_storage(num_iterations) + def _update_rays_per_sec(self, train_t: TimeWriter, vis_t: TimeWriter, step: int): + """Performs update on rays/sec calclation for training + + Args: + train_t: timer object carrying time to execute total training iteration + vis_t: timer object carrying time to execute visualization step + step: current step + """ + train_num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch + writer.put_time( + name=EventName.TRAIN_RAYS_PER_SEC, + duration=train_num_rays_per_batch / (train_t.duration - vis_t.duration), + step=step, + avg_over_steps=True, + ) + def _write_out_storage(self, step: int) -> None: """Perform writes only during appropriate time steps diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py index 15bae3c48..9e8099db9 100644 --- a/nerfactory/pipelines/base.py +++ b/nerfactory/pipelines/base.py @@ -29,8 +29,9 @@ from torch.nn.parallel import DistributedDataParallel as DDP from nerfactory.configs import base as cfg from nerfactory.datamanagers.base import DataManager from nerfactory.models.base import Model -from nerfactory.utils import profiler +from nerfactory.utils import profiler, writer from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes +from nerfactory.utils.writer import EventName, TimeWriter class Pipeline(nn.Module): @@ -122,14 +123,22 @@ class Pipeline(nn.Module): # NOTE(ethan): next_eval() is not being used right now assert self.datamanager.eval_dataloader is not None for camera_ray_bundle, batch in self.datamanager.eval_dataloader: - assert camera_ray_bundle.camera_indices is not None - image_idx = int(camera_ray_bundle.camera_indices[0, 0]) - if self.world_size > 1: - outputs = self.model.module.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - psnr = self.model.module.log_test_image_outputs(image_idx, step, batch, outputs) - else: - outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - psnr = self.model.log_test_image_outputs(image_idx, step, batch, outputs) + with TimeWriter(None, None, write=False) as test_t: + assert camera_ray_bundle.camera_indices is not None + image_idx = int(camera_ray_bundle.camera_indices[0, 0]) + if self.world_size > 1: + outputs = self.model.module.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + psnr = self.model.module.log_test_image_outputs(image_idx, step, batch, outputs) + else: + outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + psnr = self.model.log_test_image_outputs(image_idx, step, batch, outputs) + if step is not None: + writer.put_time( + name=EventName.TEST_RAYS_PER_SEC, + duration=len(camera_ray_bundle) / test_t.duration, + step=step, + avg_over_steps=True, + ) # TODO(ethan): this function should probably return something? self.train() diff --git a/nerfactory/utils/writer.py b/nerfactory/utils/writer.py index 4b568ac76..14cfd95db 100644 --- a/nerfactory/utils/writer.py +++ b/nerfactory/utils/writer.py @@ -49,7 +49,9 @@ class EventName(enum.Enum): TOTAL_TRAIN_TIME = "Train Total (time)" ITER_VIS_TIME = "Visualizer Rendering (time)" ETA = "ETA (time)" - RAYS_PER_SEC = "Rays / Sec" + TRAIN_RAYS_PER_SEC = "Train Rays / Sec" + TEST_RAYS_PER_SEC = "Test Rays / Sec" + VIS_RAYS_PER_SEC = "Vis Rays / Sec" CURR_TEST_PSNR = "Test PSNR" diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index cd8c1809e..e9d2d5e57 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/nerfactory/viewer/server/viewer_utils.py @@ -31,9 +31,9 @@ from nerfactory.cameras.rays import RayBundle from nerfactory.configs import base as cfg from nerfactory.datamanagers.datasets import InputDataset from nerfactory.models.base import Model -from nerfactory.utils import profiler, visualization +from nerfactory.utils import profiler, visualization, writer from nerfactory.utils.decorators import check_visualizer_enabled, decorate_all -from nerfactory.utils.writer import GLOBAL_BUFFER, EventName +from nerfactory.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter from nerfactory.viewer.server.subprocess import run_viewer_bridge_server_as_subprocess from nerfactory.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_h from nerfactory.viewer.server.visualizer import Viewer @@ -230,6 +230,7 @@ class VisualizerState: self.check_done_render = True self.last_render_time = time.time() self.min_wait_time = 0.5 # 1.0 is on high side and will cause lag + self.step = 0 self.outputs_set = False @@ -272,6 +273,7 @@ class VisualizerState: """ is_training = self.vis["renderingState/isTraining"].read() + self.step = step if is_training is None or is_training: # in training mode, render every few steps @@ -391,13 +393,16 @@ class VisualizerState: image = (image_output).astype("uint8") self.vis.set_image(image) - def _update_viewer_stats(self, render_time: float, image_height: int) -> None: + def _update_viewer_stats(self, render_time: float, num_rays: int, image_height: int) -> None: """Function that calculates and populates all the rendering statistics accordingly Args: render_time: total time spent rendering current view image_height: resolution of the current view """ + writer.put_time( + name=EventName.VIS_RAYS_PER_SEC, duration=num_rays / render_time, step=self.step, avg_over_steps=True + ) is_training = self.vis["renderingState/isTraining"].read() if is_training is None or is_training: # process the current rendering fps @@ -491,20 +496,18 @@ class VisualizerState: check_thread = CheckThread(state=self) render_thread = RenderThread(state=self, graph=graph, camera_ray_bundle=camera_ray_bundle) - start_time = time.time() - check_thread.start() - render_thread.start() - try: - render_thread.join() - check_thread.join() - except Exception: # pylint: disable=broad-except - pass - render_duration = time.time() - start_time + with TimeWriter(None, None, write=False) as vis_t: + check_thread.start() + render_thread.start() + try: + render_thread.join() + check_thread.join() + except Exception: # pylint: disable=broad-except + pass graph.train() - outputs = render_thread.vis_outputs if outputs is not None: stuff_colors = graph.stuff_colors if hasattr(graph, "stuff_colors") else None self._send_output_to_viewer(outputs, stuff_colors=stuff_colors) - self._update_viewer_stats(render_duration, image_height) + self._update_viewer_stats(vis_t.duration, num_rays=len(camera_ray_bundle), image_height=image_height)