moving timing around (#487)

* moved timing functions around

* lint

nerfactory/engine/trainer.py

@@ -139,9 +139,14 @@ class Trainer:
                     for callback in self.callbacks:
                         callback.run_callback_at_location(step, location=TrainingCallbackLocation.AFTER_TRAIN_ITERATION)
 
-                    vis_t = self._update_viewer_state(step)
+                writer.put_time(
+                    name=EventName.TRAIN_RAYS_PER_SEC,
+                    duration=self.config.pipeline.datamanager.train_num_rays_per_batch / train_t.duration,
+                    step=step,
+                    avg_over_steps=True,
+                )
 
-                self._update_viewer_rays_per_sec(train_t, vis_t, step)
+                self._update_viewer_state(step)
 
                 # a batch of train rays
                 if step_check(step, self.config.logging.steps_per_log, run_at_zero=True):
@@ -159,7 +164,14 @@ class Trainer:
 
                 # one eval image
                 if step_check(step, self.config.trainer.steps_per_eval_image):
-                    metrics_dict, images_dict = self.pipeline.get_eval_image_metrics_and_images(step=step)
+                    with TimeWriter(writer, EventName.TEST_RAYS_PER_SEC, write=False) as test_t:
+                        metrics_dict, images_dict = self.pipeline.get_eval_image_metrics_and_images(step=step)
+                    writer.put_time(
+                        name=EventName.TEST_RAYS_PER_SEC,
+                        duration=metrics_dict["num_rays"] / test_t.duration,
+                        step=step,
+                        avg_over_steps=True,
+                    )
                     writer.put_dict(name="Eval Images Metrics", scalar_dict=metrics_dict, step=step)
                     group = "Eval Images"
                     for image_name, image in images_dict.items():
@@ -173,7 +185,6 @@ class Trainer:
                 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)
 
-                self._update_viewer_rays_per_sec(train_t, vis_t, step)
                 self._write_out_storage(step)
 
         self._write_out_storage(num_iterations)
@@ -200,7 +211,7 @@ class Trainer:
         )
 
     @check_viewer_enabled
-    def _update_viewer_state(self, step: int) -> TimeWriter:
+    def _update_viewer_state(self, step: int):
         """Updates the viewer state by rendering out scene with current pipeline
         Returns the time taken to render scene.
 
@@ -208,7 +219,7 @@ class Trainer:
             step: current train step
         """
         assert self.viewer_state is not None
-        with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as vis_t:
+        with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as _:
             num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch
             try:
                 self.viewer_state.update_scene(step, self.pipeline.model, num_rays_per_batch)
@@ -218,7 +229,6 @@ class Trainer:
                 self.viewer_state.vis["renderingState/log_errors"].write(
                     "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing."
                 )
-        return vis_t
 
     @check_viewer_enabled
     def _update_viewer_rays_per_sec(self, train_t: TimeWriter, vis_t: TimeWriter, step: int):

nerfactory/models/instant_ngp.py

@@ -21,7 +21,7 @@ from __future__ import annotations
 from dataclasses import dataclass, field
 from typing import Dict, List, Literal, Optional, Tuple, Type
 
-import nerfacc  # pylint: disable=import-error
+import nerfacc
 import torch
 from torch.nn import Parameter
 from torchmetrics import PeakSignalNoiseRatio

nerfactory/models/modules/ray_sampler.py

@@ -20,7 +20,7 @@ import math
 from abc import abstractmethod
 from typing import Callable, Optional, Tuple
 
-import nerfacc  # pylint: disable=import-error
+import nerfacc
 import torch
 from torch import nn
 from torchtyping import TensorType

nerfactory/pipelines/base.py

@@ -30,9 +30,8 @@ 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, writer
+from nerfactory.utils import profiler
 from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes
-from nerfactory.utils.writer import EventName, TimeWriter
 
 
 def module_wrapper(module: nn.Module) -> nn.Module:
@@ -166,16 +165,12 @@ class Pipeline(nn.Module):
         """
         self.eval()
         image_idx, camera_ray_bundle, batch = self.datamanager.next_eval_image(step)
-        with TimeWriter(None, None, write=False) as test_t:
-            outputs = module_wrapper(self.model).get_outputs_for_camera_ray_bundle(camera_ray_bundle)
-            metrics_dict, images_dict = module_wrapper(self.model).get_image_metrics_and_images(outputs, batch)
-        writer.put_time(
-            name=EventName.TEST_RAYS_PER_SEC,
-            duration=len(camera_ray_bundle) / test_t.duration,
-            step=step,
-            avg_over_steps=True,
-        )
+        outputs = module_wrapper(self.model).get_outputs_for_camera_ray_bundle(camera_ray_bundle)
+        metrics_dict, images_dict = module_wrapper(self.model).get_image_metrics_and_images(outputs, batch)
+        assert "image_idx" not in metrics_dict
         metrics_dict["image_idx"] = image_idx
+        assert "num_rays" not in metrics_dict
+        metrics_dict["num_rays"] = len(camera_ray_bundle)
         self.train()
         return metrics_dict, images_dict