tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

added logic to comp rays/s on train and vis (#364)

Browse Source 
* added logic to comp rays/s on train and vis

* added test and fixed vis rays/sec
This commit is contained in:
evonneng 2022-09-07 13:43:53 -07:00 committed by GitHub
parent 7aa6ca9948
commit 7b6ed5496b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 62 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
docs/tooling/logging_profiling.md
View File

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

4
nerfactory/configs/base.py
View File

@ -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("./")

24
nerfactory/engine/trainer.py
View File

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

27
nerfactory/pipelines/base.py
View File

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

4
nerfactory/utils/writer.py
View File

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

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

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