profiling

nerfstudio/cameras/cameras.py

@@ -35,6 +35,7 @@ from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.data.scene_box import SceneBox
 from nerfstudio.utils.misc import strtobool
 from nerfstudio.utils.tensor_dataclass import TensorDataclass
+from nerfstudio.utils import profiler
 
 TORCH_DEVICE = Union[torch.device, str]  # pylint: disable=invalid-name
 
@@ -309,6 +310,8 @@ class Cameras(TensorDataclass):
             image_coords = torch.stack(image_coords, dim=-1) + pixel_offset  # stored as (y, x) coordinates
         return image_coords
 
+    @profiler.time_function
+    @profile
     def generate_rays(  # pylint: disable=too-many-statements
         self,
         camera_indices: Union[TensorType["num_rays":..., "num_cameras_batch_dims"], int],
@@ -487,6 +490,7 @@ class Cameras(TensorDataclass):
         return raybundle
 
     # pylint: disable=too-many-statements
+    @profile
     def _generate_rays_from_coords(
         self,
         camera_indices: TensorType["num_rays":..., "num_cameras_batch_dims"],

nerfstudio/data/datamanagers/base_datamanager.py

@@ -71,6 +71,7 @@ from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate
 from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes
 from nerfstudio.model_components.ray_generators import RayGenerator
 from nerfstudio.utils.misc import IterableWrapper
+from nerfstudio.utils import profiler
 
 CONSOLE = Console(width=120)
 
@@ -510,6 +511,8 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
             num_workers=self.world_size * 4,
         )
 
+    @profiler.time_function
+    @profile
     def next_train(self, step: int) -> Tuple[RayBundle, Dict]:
         """Returns the next batch of data from the train dataloader."""
         self.train_count += 1

nerfstudio/data/pixel_samplers.py

@@ -21,6 +21,7 @@ from typing import Dict, Optional, Union
 
 import torch
 from torchtyping import TensorType
+from nerfstudio.utils import profiler
 
 
 class PixelSampler:  # pylint: disable=too-few-public-methods
@@ -182,6 +183,7 @@ class PixelSampler:  # pylint: disable=too-few-public-methods
 
         return collated_batch
 
+    @profiler.time_function
     def sample(self, image_batch: Dict):
         """Sample an image batch and return a pixel batch.
 

nerfstudio/model_components/ray_generators.py

@@ -21,6 +21,7 @@ from torchtyping import TensorType
 from nerfstudio.cameras.camera_optimizers import CameraOptimizer
 from nerfstudio.cameras.cameras import Cameras
 from nerfstudio.cameras.rays import RayBundle
+from nerfstudio.utils import profiler
 
 
 class RayGenerator(nn.Module):
@@ -38,6 +39,8 @@ class RayGenerator(nn.Module):
         self.pose_optimizer = pose_optimizer
         self.register_buffer("image_coords", cameras.get_image_coords(), persistent=False)
 
+    @profiler.time_function
+    @profile
     def forward(self, ray_indices: TensorType["num_rays", 3]) -> RayBundle:
         """Index into the cameras to generate the rays.
 

nerfstudio/models/base_model.py

@@ -33,6 +33,7 @@ from nerfstudio.configs.config_utils import to_immutable_dict
 from nerfstudio.data.scene_box import SceneBox
 from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes
 from nerfstudio.model_components.scene_colliders import NearFarCollider
+from nerfstudio.utils import profiler
 
 
 # Model related configs
@@ -126,6 +127,7 @@ class Model(nn.Module):
             Outputs of model. (ie. rendered colors)
         """
 
+    @profiler.time_function
     def forward(self, ray_bundle: RayBundle) -> Dict[str, torch.Tensor]:
         """Run forward starting with a ray bundle. This outputs different things depending on the configuration
         of the model and whether or not the batch is provided (whether or not we are training basically)

nerfstudio/models/nerfacto.py

@@ -60,6 +60,7 @@ from nerfstudio.model_components.scene_colliders import NearFarCollider
 from nerfstudio.model_components.shaders import NormalsShader
 from nerfstudio.models.base_model import Model, ModelConfig
 from nerfstudio.utils import colormaps
+from nerfstudio.utils import profiler
 
 
 @dataclass
@@ -262,6 +263,7 @@ class NerfactoModel(Model):
             )
         return callbacks
 
+    @profiler.time_function
     def get_outputs(self, ray_bundle: RayBundle):
         ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns)
         field_outputs = self.field(ray_samples, compute_normals=self.config.predict_normals)
@@ -305,6 +307,7 @@ class NerfactoModel(Model):
 
         return outputs
 
+    @profiler.time_function
     def get_metrics_dict(self, outputs, batch):
         metrics_dict = {}
         image = batch["image"].to(self.device)
@@ -313,6 +316,7 @@ class NerfactoModel(Model):
             metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"])
         return metrics_dict
 
+    @profiler.time_function
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
         loss_dict = {}
         image = batch["image"].to(self.device)

nerfstudio/pipelines/base_pipeline.py

@@ -266,6 +266,7 @@ class VanillaPipeline(Pipeline):
         return self.model.device
 
     @profiler.time_function
+    @profile
     def get_train_loss_dict(self, step: int):
         """This function gets your training loss dict. This will be responsible for
         getting the next batch of data from the DataManager and interfacing with the