set pose_optimizer to None if it's mode is "off"

nerfstudio/data/datamanagers/base_datamanager.py

@@ -70,8 +70,8 @@ from nerfstudio.data.utils.dataloaders import (
 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
+from nerfstudio.utils.misc import IterableWrapper
 
 CONSOLE = Console(width=120)
 
@@ -469,9 +469,12 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         )
         self.iter_train_image_dataloader = iter(self.train_image_dataloader)
         self.train_pixel_sampler = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch)
-        self.train_camera_optimizer = self.config.camera_optimizer.setup(
-            num_cameras=self.train_dataset.cameras.size, device=self.device
-        )
+        if self.config.camera_optimizer.mode != "off":
+            self.train_camera_optimizer = self.config.camera_optimizer.setup(
+                num_cameras=self.train_dataset.cameras.size, device=self.device
+            )
+        else:
+            self.train_camera_optimizer = None
         self.train_ray_generator = RayGenerator(
             self.train_dataset.cameras.to(self.device),
             self.train_camera_optimizer,
@@ -492,9 +495,12 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         )
         self.iter_eval_image_dataloader = iter(self.eval_image_dataloader)
         self.eval_pixel_sampler = self._get_pixel_sampler(self.eval_dataset, self.config.eval_num_rays_per_batch)
-        self.eval_camera_optimizer = self.config.camera_optimizer.setup(
-            num_cameras=self.eval_dataset.cameras.size, device=self.device
-        )
+        if self.config.camera_optimizer.mode != "off":
+            self.eval_camera_optimizer = self.config.camera_optimizer.setup(
+                num_cameras=self.eval_dataset.cameras.size, device=self.device
+            )
+        else:
+            self.eval_camera_optimizer = None
         self.eval_ray_generator = RayGenerator(
             self.eval_dataset.cameras.to(self.device),
             self.eval_camera_optimizer,
@@ -556,11 +562,9 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         """
         param_groups = {}
 
-        camera_opt_params = list(self.train_camera_optimizer.parameters())
         if self.config.camera_optimizer.mode != "off":
+            camera_opt_params = list(self.train_camera_optimizer.parameters())
             assert len(camera_opt_params) > 0
             param_groups[self.config.camera_optimizer.param_group] = camera_opt_params
-        else:
-            assert len(camera_opt_params) == 0
 
         return param_groups

nerfstudio/data/pixel_samplers.py

@@ -74,6 +74,7 @@ class PixelSampler:  # pylint: disable=too-few-public-methods
 
         return indices
 
+    @profile
     def collate_image_dataset_batch(self, batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False):
         """
         Operates on a batch of images and samples pixels to use for generating rays.
@@ -112,6 +113,7 @@ class PixelSampler:  # pylint: disable=too-few-public-methods
 
         return collated_batch
 
+    @profile
     def collate_image_dataset_batch_list(self, batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False):
         """
         Does the same as collate_image_dataset_batch, except it will operate over a list of images / masks inside
@@ -184,6 +186,7 @@ class PixelSampler:  # pylint: disable=too-few-public-methods
         return collated_batch
 
     @profiler.time_function
+    @profile
     def sample(self, image_batch: Dict):
         """Sample an image batch and return a pixel batch.
 

nerfstudio/model_components/ray_generators.py

@@ -15,6 +15,9 @@
 """
 Ray generator.
 """
+from typing import Optional
+
+import torch
 from torch import nn
 from torchtyping import TensorType
 
@@ -33,7 +36,11 @@ class RayGenerator(nn.Module):
         pose_optimizer: pose optimization module, for optimizing noisy camera intrinsics/extrinsics.
     """
 
-    def __init__(self, cameras: Cameras, pose_optimizer: CameraOptimizer) -> None:
+    image_coords: TensorType["height", "width", 2]
+
+    def __init__(
+        self, cameras: Cameras, pose_optimizer: Optional[CameraOptimizer] = None
+    ) -> None:
         super().__init__()
         self.cameras = cameras
         self.pose_optimizer = pose_optimizer
@@ -52,7 +59,10 @@ class RayGenerator(nn.Module):
         x = ray_indices[:, 2]  # col indices
         coords = self.image_coords[y, x]
 
-        camera_opt_to_camera = self.pose_optimizer(c)
+        if self.pose_optimizer is not None:
+            camera_opt_to_camera = self.pose_optimizer(c)
+        else:
+            camera_opt_to_camera = None
 
         ray_bundle = self.cameras.generate_rays(
             camera_indices=c.unsqueeze(-1),