From 42501619229664f5949e7400f28ba0fa82589596 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Mon, 19 Sep 2022 15:58:40 -0700 Subject: [PATCH] moving timing around (#487) * moved timing functions around * lint --- nerfactory/engine/trainer.py | 24 +++++++++++++++++------- nerfactory/models/instant_ngp.py | 2 +- nerfactory/models/modules/ray_sampler.py | 2 +- nerfactory/pipelines/base.py | 17 ++++++----------- 4 files changed, 25 insertions(+), 20 deletions(-) diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index 65aa25b89..83b5de593 100644 --- a/nerfactory/engine/trainer.py +++ b/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): diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index 4e36be1e0..4c7eee8f1 100644 --- a/nerfactory/models/instant_ngp.py +++ b/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 diff --git a/nerfactory/models/modules/ray_sampler.py b/nerfactory/models/modules/ray_sampler.py index 9993c9e9d..929d06b40 100644 --- a/nerfactory/models/modules/ray_sampler.py +++ b/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 diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py index 29d4f209f..c5afa5314 100644 --- a/nerfactory/pipelines/base.py +++ b/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