diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index 03b34f250..83b3d6fee 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -102,8 +102,13 @@ class Trainer: num_iterations = self.config.trainer.max_num_iterations for step in range(self.start_step, self.start_step + num_iterations): + # Note: if 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: loss_metric_dict = self.train_iteration(step) + + with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as t: + self.visualizer_state.update_scene(step, self.pipeline.model) + train_num_rays_per_batch = self.pipeline.dataloader.train_num_rays_per_batch writer.put_scalar(name=EventName.RAYS_PER_SEC, scalar=train_num_rays_per_batch / t.duration, step=step) @@ -114,7 +119,6 @@ class Trainer: if step % self.config.trainer.steps_per_test == 0: self.pipeline.get_eval_loss_dict(step=step) self._write_out_storage(step) - self.visualizer_state.update_scene(step, self.pipeline.model) self._write_out_storage(num_iterations) diff --git a/nerfactory/utils/writer.py b/nerfactory/utils/writer.py index 810118fa8..ae9695b80 100644 --- a/nerfactory/utils/writer.py +++ b/nerfactory/utils/writer.py @@ -49,6 +49,7 @@ class EventName(enum.Enum): ITER_LOAD_TIME = "Data Load (time)" ITER_TRAIN_TIME = "Train Iter (time)" TOTAL_TRAIN_TIME = "Train Total (time)" + ITER_VIS_TIME = "Visualizer Rendering (time)" ETA = "ETA (time)" RAYS_PER_SEC = "Rays / Sec" CURR_TEST_PSNR = "Test PSNR" @@ -116,7 +117,7 @@ def put_time(name: str, duration: float, step: int, avg_over_steps: bool = True, remain_iter = GLOBAL_BUFFER["max_iter"] - step remain_time = remain_iter * GLOBAL_BUFFER["events"][name]["avg"] put_scalar(EventName.ETA, remain_time, step) - GLOBAL_BUFFER["viewer/train_eta"] = _format_time(remain_time) + GLOBAL_BUFFER["events"][EventName.ETA.value] = _format_time(remain_time) @check_main_thread @@ -153,7 +154,6 @@ def setup_event_writers(config: LoggingConfig, max_iter: int) -> None: GLOBAL_BUFFER["max_buffer_size"] = config.max_buffer_size GLOBAL_BUFFER["steps_per_log"] = config.steps_per_log GLOBAL_BUFFER["events"] = {} - GLOBAL_BUFFER["viewer/train_eta"] = "Paused" class Writer: diff --git a/nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx b/nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx index 91ee4bade..841aa68c3 100644 --- a/nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx +++ b/nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx @@ -33,6 +33,7 @@ export function RenderControls() { ); let eval_fps = useSelector((state) => state.renderingState.eval_fps); let train_eta = useSelector((state) => state.renderingState.train_eta); + let vis_train_ratio = useSelector((state) => state.renderingState.vis_train_ratio); const dispatch = useDispatch(); @@ -244,8 +245,13 @@ export function RenderControls() { value: train_eta, disabled: true, }, + vis_train_ratio: { + label: 'Time Allocation', + value: vis_train_ratio, + disabled: true, + }, }), - [isWebsocketConnected, isWebrtcConnected, eval_fps, train_eta], + [isWebsocketConnected, isWebrtcConnected, eval_fps, train_eta, vis_train_ratio], ); useEffect(() => { @@ -271,14 +277,16 @@ export function RenderControls() { const cmd = msgpack.decode(new Uint8Array(originalCmd.data)); if (cmd.path === '/renderingState/eval_fps') { eval_fps = cmd.data; - console.log('setting state', eval_fps); setState({ eval_fps }); } if (cmd.path === '/renderingState/train_eta') { train_eta = cmd.data; - console.log('setting state', train_eta); setState({ train_eta }); } + if (cmd.path === '/renderingState/vis_train_ratio') { + vis_train_ratio = cmd.data; + setState({ vis_train_ratio }); + } }); }, []); diff --git a/nerfactory/viewer/app/src/reducer.js b/nerfactory/viewer/app/src/reducer.js index ae3b0bab6..7d084abe6 100644 --- a/nerfactory/viewer/app/src/reducer.js +++ b/nerfactory/viewer/app/src/reducer.js @@ -20,6 +20,7 @@ const initialState = { field_of_view: 80, eval_fps: 'Paused', train_eta: 'Paused', + vis_train_ratio: 'Paused', }, // the scene state sceneState: { diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index 6b9ec79ba..21a2006ce 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/nerfactory/viewer/server/viewer_utils.py @@ -32,7 +32,7 @@ from nerfactory.models.base import Model from nerfactory.utils import profiler from nerfactory.utils.config import ViewerConfig from nerfactory.utils.decorators import check_visualizer_enabled, decorate_all -from nerfactory.utils.writer import GLOBAL_BUFFER +from nerfactory.utils.writer import GLOBAL_BUFFER, EventName from nerfactory.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_h from nerfactory.viewer.server.visualizer import Viewer @@ -280,15 +280,32 @@ class VisualizerState: image = (image_output).astype("uint8") self.vis.set_image(image) - def _update_viewer_stats(self, render_time, image_height) -> None: + def _update_viewer_stats(self, render_time: float, 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 + """ is_training = self.vis["renderingState/isTraining"].read() if is_training is None or is_training: + # process the current rendering fps eval_fps = f"{1 / render_time:.2f} fps at {image_height} res" self.vis["renderingState/eval_fps"].write(eval_fps) - self.vis["renderingState/train_eta"].write(GLOBAL_BUFFER["viewer/train_eta"]) + # 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"]: + 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" + self.vis["renderingState/vis_train_ratio"].write(vis_train_ratio) + else: + self.vis["renderingState/vis_train_ratio"] = "Starting" else: self.vis["renderingState/eval_fps"].write("Paused") self.vis["renderingState/train_eta"].write("Paused") + self.vis["renderingState/vis_train_ratio"].write("100% spent on viewer") @profiler.time_function def _render_image_in_viewer(self, graph: Model) -> None: