From 75d2d14b75380c9843461ec3dfd69db923fab453 Mon Sep 17 00:00:00 2001 From: Zhuoyang Date: Tue, 19 Sep 2023 19:47:50 -0700 Subject: [PATCH] add number of gaussians in logger; allow testing in nerfacto --- nerfstudio/configs/base_config.py | 1 + nerfstudio/engine/trainer.py | 24 ++++++++++++------- nerfstudio/models/gaussian_splatting.py | 4 ++++ nerfstudio/utils/writer.py | 3 ++- .../viewer/server/render_state_machine.py | 9 ++++--- 5 files changed, 29 insertions(+), 12 deletions(-) diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index 315659acb..37958dd4a 100644 --- a/nerfstudio/configs/base_config.py +++ b/nerfstudio/configs/base_config.py @@ -88,6 +88,7 @@ class LocalWriterConfig(InstantiateConfig): writer.EventName.VIS_RAYS_PER_SEC, writer.EventName.TEST_RAYS_PER_SEC, writer.EventName.ETA, + writer.EventName.GAUSSIAN_NUM, ) """specifies which stats will be logged/printed to terminal""" max_log_size: int = 10 diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 94256809c..fcce9bffb 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -36,6 +36,7 @@ from nerfstudio.configs.experiment_config import ExperimentConfig from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation from nerfstudio.engine.optimizers import Optimizers +from nerfstudio.models.gaussian_splatting import GaussianSplattingModel from nerfstudio.pipelines.base_pipeline import VanillaPipeline from nerfstudio.utils import profiler, writer from nerfstudio.utils.decorators import check_eval_enabled, check_main_thread, check_viewer_enabled @@ -262,14 +263,21 @@ class Trainer: # Skip the first two steps to avoid skewed timings that break the viewer rendering speed estimate. if step > 1: - writer.put_time( - name=EventName.TRAIN_RAYS_PER_SEC, - duration=self.world_size - * self.pipeline.datamanager.get_train_rays_per_batch() - / max(0.001, train_t.duration), - step=step, - avg_over_steps=True, - ) + if isinstance(self.pipeline.model, GaussianSplattingModel): + writer.put_scalar( + name=EventName.GAUSSIAN_NUM, + scalar=self.pipeline.model.get_gaussian_num, + step=step, + ) + else: + writer.put_time( + name=EventName.TRAIN_RAYS_PER_SEC, + duration=self.world_size + * self.pipeline.datamanager.get_train_rays_per_batch() + / max(0.001, train_t.duration), + step=step, + avg_over_steps=True, + ) self._update_viewer_state(step) diff --git a/nerfstudio/models/gaussian_splatting.py b/nerfstudio/models/gaussian_splatting.py index d9bb48db5..677abcc84 100644 --- a/nerfstudio/models/gaussian_splatting.py +++ b/nerfstudio/models/gaussian_splatting.py @@ -191,6 +191,10 @@ class GaussianSplattingModel(Model): def get_rotation(self): return self.rotation_activation(self._rotation) + @property + def get_gaussian_num(self): + return self._xyz.shape[0] + @property def get_xyz(self): return self._xyz diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index 8d3b07b1a..2b6f37f51 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -56,6 +56,7 @@ class EventName(enum.Enum): TRAIN_RAYS_PER_SEC = "Train Rays / Sec" TEST_RAYS_PER_SEC = "Test Rays / Sec" VIS_RAYS_PER_SEC = "Vis Rays / Sec" + GAUSSIAN_NUM = "Number of Gaussians" CURR_TEST_PSNR = "Test PSNR" @@ -468,7 +469,7 @@ class LocalWriter: if name in self.stats_to_track: if "(time)" in name: v = _format_time(v) - elif "Rays" in name: + elif "Rays" in name or "Gaussians" in name: v = human_format(v) else: v = f"{v:0.4f}" diff --git a/nerfstudio/viewer/server/render_state_machine.py b/nerfstudio/viewer/server/render_state_machine.py index 2a01b1a47..b02dfae9d 100644 --- a/nerfstudio/viewer/server/render_state_machine.py +++ b/nerfstudio/viewer/server/render_state_machine.py @@ -25,6 +25,7 @@ import torch from nerfstudio.cameras.cameras import Cameras from nerfstudio.model_components.renderers import \ background_color_override_context +from nerfstudio.models.gaussian_splatting import GaussianSplattingModel from nerfstudio.utils import colormaps, writer from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter from nerfstudio.viewer.server import viewer_utils @@ -129,7 +130,9 @@ class RenderStateMachine(threading.Thread): with self.viewer.train_lock if self.viewer.train_lock is not None else contextlib.nullcontext(): #TODO jake-austin: Make this check whether the model inherits from a camera based model or a ray based model - if True: + #TODO Zhuoyang: First made some dummy judgements, need to be fixed later + isGaussianSplattingModel = isinstance(self.viewer.get_model(), GaussianSplattingModel) + if isGaussianSplattingModel: pass else: camera_ray_bundle = camera.generate_rays(camera_indices=0, aabb_box=self.viewer.get_model().render_aabb) @@ -147,13 +150,13 @@ class RenderStateMachine(threading.Thread): device=self.viewer.get_model().device, ) with background_color_override_context(background_color), torch.no_grad(): - if True: + if isGaussianSplattingModel: outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera) else: outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle) else: with torch.no_grad(): - if True: + if isGaussianSplattingModel: outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera) else: outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle)