diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index e0ed08e4f..244eb3fe0 100644 --- a/nerfstudio/cameras/cameras.py +++ b/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"], diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 7e2ebedeb..04d02b679 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/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 diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 244498df8..1d5a9697e 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/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. diff --git a/nerfstudio/model_components/ray_generators.py b/nerfstudio/model_components/ray_generators.py index a894da854..2e568ee3c 100644 --- a/nerfstudio/model_components/ray_generators.py +++ b/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. diff --git a/nerfstudio/models/base_model.py b/nerfstudio/models/base_model.py index a73ea67a1..637e7a844 100644 --- a/nerfstudio/models/base_model.py +++ b/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) diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 53173640a..e21f7c8de 100644 --- a/nerfstudio/models/nerfacto.py +++ b/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) diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index c572bebeb..c1dccc77f 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/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 diff --git a/train.py.lprof b/train.py.lprof new file mode 100644 index 000000000..8fbdf9351 Binary files /dev/null and b/train.py.lprof differ