tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

added stats- ratio train:visualizer (#268)

Browse Source 
added stats: ratio train:visualizer
This commit is contained in:
evonneng 2022-08-11 16:34:57 -07:00 committed by GitHub
parent 3cb44f3a8e
commit 1499fd1cd7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 39 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
nerfactory/engine/trainer.py
View File

@ -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)

4
nerfactory/utils/writer.py
View File

@ -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:

14
nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx
View File

@ -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 });
}
});
}, []);

1
nerfactory/viewer/app/src/reducer.js
View File

@ -20,6 +20,7 @@ const initialState = {
field_of_view: 80,
eval_fps: 'Paused',
train_eta: 'Paused',
vis_train_ratio: 'Paused',
},
// the scene state
sceneState: {

23
nerfactory/viewer/server/viewer_utils.py
View File

@ -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: