From 07698dfb71e1bb334dec0bc9a2913efda32761e9 Mon Sep 17 00:00:00 2001 From: Ruilong Li Date: Thu, 6 Apr 2023 20:30:56 +0000 Subject: [PATCH] profiling --- nerfstudio/cameras/cameras.py | 4 ++++ nerfstudio/data/datamanagers/base_datamanager.py | 3 +++ nerfstudio/data/pixel_samplers.py | 2 ++ nerfstudio/model_components/ray_generators.py | 3 +++ nerfstudio/models/base_model.py | 2 ++ nerfstudio/models/nerfacto.py | 4 ++++ nerfstudio/pipelines/base_pipeline.py | 1 + train.py.lprof | Bin 0 -> 2380 bytes 8 files changed, 19 insertions(+) create mode 100644 train.py.lprof 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 0000000000000000000000000000000000000000..8fbdf9351162b03a3213724fa0ba25018812537f GIT binary patch literal 2380 zcmb8xdrVVT90zd9tfjmad>|?T9CeCxqEkUYWfP)pRYY(OwYk#U_G;T(@4XeuV%vbr}sSW z;dj`qZCRiX!IwJTC3twNR}$@liSR0(I>J-5(cpNg2 zZ*+=o-YEG5m*{c01Y?CLZIHcOIdAmvl3iAOHbH#0Sght@V6HKQG*x-z(WG8Ua z9!+QkZtM>u><8wx^(HeM1un)%Qt}lL`dNa8;XELZz|$@OzPs3iZXmaCGJWQ2;KNQ# z^b#SUJ-A}`3@ytpBl|ZS!(Iin1yS^=?|~08k^Bh+b_`<)@6=iT)~MsG&%Ft&T@u~a za#56QPv1&TD7}@`=59TETrNC7cc|$70 z7T`hhz-Weg;MG+)&kcb6NEXdb1v=7>k7C#k%)E4$k{!U-wjz36O+fQLZ8XD9U|H)o zG@)i-OaJkFhWCJdNrm$mb^(9%?WJxlgQ!fR2OI=u+LTC!c3{?{bo%o;2CV8c^)RA^QO^UrrBa$#pgUVsblW#jPb@)x}<+@g(lvFhXYiXPsh$VPTKJgrZnB-7<#nxg9L zqEx|2w!vG&lu6~O?4LtyNYe~zcimh@?U~emY&KE*HEQ=PsbVj{G@G>p+O`9C^4RGx zWkH&{`(LD^66PRnUTmIq5$*i;ul?vSMc)mAH>iio)RmF)TnU|erN$! zmnh3tn^3L{EoC&_L;-DqSYrue^B6P8NnGpubPmT3a?!^jYiR}JaBY=YGHKn3K5t U9zh8f=8edR=uf-R;#-&TKdp>sJ^%m! literal 0 HcmV?d00001