initial impl of parallel data manager

nerfstudio/configs/method_configs.py

@@ -27,6 +27,7 @@ from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
 from nerfstudio.configs.base_config import ViewerConfig
 from nerfstudio.configs.external_methods import get_external_methods
 from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager, VanillaDataManagerConfig
+from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManagerConfig
 from nerfstudio.data.datamanagers.random_cameras_datamanager import RandomCamerasDataManagerConfig
 from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
 from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig
@@ -86,7 +87,8 @@ method_configs["nerfacto"] = TrainerConfig(
     max_num_iterations=30000,
     mixed_precision=True,
     pipeline=VanillaPipelineConfig(
-        datamanager=VanillaDataManagerConfig(
+        datamanager=ParallelDataManagerConfig(
+        # datamanager=VanillaDataManagerConfig(
             dataparser=NerfstudioDataParserConfig(),
             train_num_rays_per_batch=4096,
             eval_num_rays_per_batch=4096,
@@ -107,7 +109,7 @@ method_configs["nerfacto"] = TrainerConfig(
         },
         "camera_opt": {
             "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
-            "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=2000),
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=5000),
         },
     },
     viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
@@ -199,12 +201,12 @@ method_configs["nerfacto-huge"] = TrainerConfig(
             "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=50000),
         },
         "camera_opt": {
-            "optimizer": RAdamOptimizerConfig(lr=6e-4, eps=1e-8),
-            "scheduler": ExponentialDecaySchedulerConfig(lr_final=6e-5, max_steps=50000),
+            "optimizer": RAdamOptimizerConfig(lr=1e-3, eps=1e-8),
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=10000),
         },
     },
     viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
-    vis="viewer",
+    vis="viewer_beta",
 )
 
 # method_configs["depth-nerfacto"] = TrainerConfig(

nerfstudio/data/datamanagers/base_datamanager.py

@@ -23,19 +23,7 @@ from abc import abstractmethod
 from collections import defaultdict
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Generic,
-    List,
-    Literal,
-    Optional,
-    Tuple,
-    Type,
-    Union,
-    cast,
-)
+from typing import Any, Callable, Dict, Generic, List, Literal, Optional, Tuple, Type, Union, cast
 
 import torch
 from torch import nn
@@ -50,16 +38,8 @@ from nerfstudio.configs.dataparser_configs import AnnotatedDataParserUnion
 from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs
 from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
 from nerfstudio.data.datasets.base_dataset import InputDataset
-from nerfstudio.data.pixel_samplers import (
-    EquirectangularPixelSampler,
-    PatchPixelSampler,
-    PixelSampler,
-)
-from nerfstudio.data.utils.dataloaders import (
-    CacheDataloader,
-    FixedIndicesEvalDataloader,
-    RandIndicesEvalDataloader,
-)
+from nerfstudio.data.pixel_samplers import EquirectangularPixelSampler, PatchPixelSampler, PixelSampler
+from nerfstudio.data.utils.dataloaders import CacheDataloader, FixedIndicesEvalDataloader, RandIndicesEvalDataloader
 from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate
 from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes
 from nerfstudio.model_components.ray_generators import RayGenerator

nerfstudio/data/datamanagers/parallel_datamanager.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import concurrent.futures
+import functools
+import torch.multiprocessing as mp
+import time
+from abc import abstractmethod
+from collections import defaultdict
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import (Any, Callable, Dict, Generic, List, Literal, Optional,
+                    Tuple, Type, Union, cast)
+
+import torch
+from rich.progress import track
+from torch import nn
+from torch.nn import Parameter
+from torch.utils.data.distributed import DistributedSampler
+from typing_extensions import TypeVar
+
+from nerfstudio.cameras.cameras import CameraType
+from nerfstudio.cameras.rays import RayBundle
+from nerfstudio.configs.base_config import InstantiateConfig
+from nerfstudio.configs.dataparser_configs import AnnotatedDataParserUnion
+from nerfstudio.data.datamanagers.base_datamanager import (
+    DataManager, DataManagerConfig, TDataset, variable_res_collate)
+from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs
+from nerfstudio.data.dataparsers.blender_dataparser import \
+    BlenderDataParserConfig
+from nerfstudio.data.datasets.base_dataset import InputDataset
+from nerfstudio.data.pixel_samplers import (EquirectangularPixelSampler,
+                                            PatchPixelSampler, PixelSampler)
+from nerfstudio.data.utils.dataloaders import (CacheDataloader,
+                                               FixedIndicesEvalDataloader,
+                                               RandIndicesEvalDataloader)
+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.rich_utils import CONSOLE
+
+
+@dataclass
+class ParallelDataManagerConfig(DataManagerConfig):
+    """A basic data manager"""
+
+    _target: Type = field(default_factory=lambda: ParallelDataManager)
+    """Target class to instantiate."""
+    dataparser: AnnotatedDataParserUnion = BlenderDataParserConfig()
+    """Specifies the dataparser used to unpack the data."""
+    train_num_rays_per_batch: int = 1024
+    """Number of rays per batch to use per training iteration."""
+    train_num_images_to_sample_from: int = -1
+    """Number of images to sample during training iteration."""
+    train_num_times_to_repeat_images: int = -1
+    """When not training on all images, number of iterations before picking new
+    images. If -1, never pick new images."""
+    eval_num_rays_per_batch: int = 1024
+    """Number of rays per batch to use per eval iteration."""
+    eval_num_images_to_sample_from: int = -1
+    """Number of images to sample during eval iteration."""
+    eval_num_times_to_repeat_images: int = -1
+    """When not evaluating on all images, number of iterations before picking
+    new images. If -1, never pick new images."""
+    eval_image_indices: Optional[Tuple[int, ...]] = (0,)
+    """Specifies the image indices to use during eval; if None, uses all."""
+    collate_fn: Callable[[Any], Any] = cast(Any, staticmethod(nerfstudio_collate))
+    """Specifies the collate function to use for the train and eval dataloaders."""
+    camera_res_scale_factor: float = 1.0
+    """The scale factor for scaling spatial data such as images, mask, semantics
+    along with relevant information about camera intrinsics
+    """
+    patch_size: int = 1
+    """Size of patch to sample from. If >1, patch-based sampling will be used."""
+    n_procs: int = 1
+    """Number of processes to use for data loading."""
+
+class DataProc(mp.Process):
+    def __init__(
+        self,
+        out_queue: mp.Queue,
+        config: ParallelDataManagerConfig,
+        dataparser_outputs: DataparserOutputs,
+        train_dataset: TDataset,
+        pix_sampler: PixelSampler,
+    ):
+        super().__init__()
+        self.daemon = True
+        self.out_queue = out_queue
+        self.config = config
+        self.dataparser_outputs = dataparser_outputs
+        self.train_dataset = train_dataset
+        self.exclude_batch_keys_from_device = self.train_dataset.exclude_batch_keys_from_device
+        if self.dataparser_outputs is not None:
+            cameras = self.dataparser_outputs.cameras
+            if len(cameras) > 1:
+                for i in range(1, len(cameras)):
+                    if cameras[0].width != cameras[i].width or cameras[0].height != cameras[i].height:
+                        CONSOLE.print("Variable resolution, using variable_res_collate")
+                        self.config.collate_fn = variable_res_collate
+                        break
+        self.train_pixel_sampler = pix_sampler
+        self.train_ray_generator = RayGenerator(self.train_dataset.cameras)
+        self.cache_images()
+    
+    def run(self):
+        while True:
+            batch = self.train_pixel_sampler.sample(self.img_data)
+            ray_indices = batch["indices"]
+            ray_bundle:RayBundle = self.train_ray_generator(ray_indices)
+            ray_bundle = ray_bundle.pin_memory()
+            self.out_queue.put((ray_bundle,batch), block=True)
+
+    def cache_images(self):
+        # caches all the input images in a NxHxWx3 tensor
+        indices = range(len(self.train_dataset))
+        batch_list = []
+        results = []
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            for idx in indices:
+                res = executor.submit(self.train_dataset.__getitem__, idx)
+                results.append(res)
+            for res in track(results, description="Loading data batch", transient=True):
+                batch_list.append(res.result())
+        self.img_data = self.config.collate_fn(batch_list)
+
+
+class ParallelDataManager(DataManager, Generic[TDataset]):
+    def __init__(
+        self,
+        config: ParallelDataManagerConfig,
+        device: Union[torch.device, str] = "cpu",
+        test_mode: Literal["test", "val", "inference"] = "val",
+        world_size: int = 1,
+        local_rank: int = 0,
+        **kwargs,
+    ):
+        self.dataset_type: Type[TDataset] = kwargs.get("_dataset_type", getattr(TDataset, "__default__"))
+        self.config = config
+        self.device = device
+        self.world_size = world_size
+        self.local_rank = local_rank
+        # self.sampler = None
+        self.test_mode = test_mode
+        self.test_split = "test" if test_mode in ["test", "inference"] else "val"
+        self.dataparser_config = self.config.dataparser
+        if self.config.data is not None:
+            self.config.dataparser.data = Path(self.config.data)
+        else:
+            self.config.data = self.config.dataparser.data
+        self.dataparser = self.dataparser_config.setup()
+        if test_mode == "inference":
+            self.dataparser.downscale_factor = 1  # Avoid opening images
+        self.includes_time = self.dataparser.includes_time
+        self.train_dataparser_outputs: DataparserOutputs = self.dataparser.get_dataparser_outputs(split="train")
+        self.train_dataset = self.create_train_dataset()
+        # self.eval_dataset = self.create_eval_dataset()
+        self.pix_sampler = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch)
+        mp.set_start_method('spawn')
+        self.data_q = mp.Queue(maxsize=4)
+        self.data_procs = [
+            DataProc(self.data_q, self.config, self.train_dataparser_outputs, self.train_dataset, self.pix_sampler)
+            for i in range(self.config.n_procs)
+        ]
+        for proc in self.data_procs:
+            proc.start()
+        super().__init__()
+        # Prime the executor with the first batch
+        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+        self.batch_fut = self.executor.submit(self.data_q.get)
+
+    def create_train_dataset(self) -> TDataset:
+        """Sets up the data loaders for training"""
+        return self.dataset_type(
+            dataparser_outputs=self.train_dataparser_outputs,
+            scale_factor=self.config.camera_res_scale_factor,
+        )
+
+    def _get_pixel_sampler(self, dataset: TDataset, *args: Any, **kwargs: Any) -> PixelSampler:
+        """Infer pixel sampler to use."""
+        if self.config.patch_size > 1:
+            return PatchPixelSampler(*args, **kwargs, patch_size=self.config.patch_size)
+
+        # If all images are equirectangular, use equirectangular pixel sampler
+        is_equirectangular = dataset.cameras.camera_type == CameraType.EQUIRECTANGULAR.value
+        if is_equirectangular.all():
+            return EquirectangularPixelSampler(*args, **kwargs)
+        # Otherwise, use the default pixel sampler
+        if is_equirectangular.any():
+            CONSOLE.print("[bold yellow]Warning: Some cameras are equirectangular, but using default pixel sampler.")
+        return PixelSampler(*args, **kwargs)
+    
+    def next_train(self, step: int) -> Tuple[RayBundle, Dict]:
+        self.train_count += 1
+        # Fetch the next batch in an executor to parallelize the queue get() operation
+        # with the train step
+        bundle, batch = self.batch_fut.result()
+        self.batch_fut = self.executor.submit(self.data_q.get)
+        bundle = bundle.to(self.device)
+        return bundle,batch
+    
+    def get_train_rays_per_batch(self) -> int:
+        return self.config.train_num_rays_per_batch
+
+    def get_eval_rays_per_batch(self) -> int:
+        return self.config.eval_num_rays_per_batch
+
+    def get_datapath(self) -> Path:
+        return self.config.dataparser.data
+
+    def get_param_groups(self) -> Dict[str, List[Parameter]]:
+        """Get the param groups for the data manager.
+        Returns:
+            A list of dictionaries containing the data manager's param groups.
+        """
+        return {}

nerfstudio/data/pixel_samplers.py

@@ -67,10 +67,10 @@ class PixelSampler:
             chosen_indices = random.sample(range(len(nonzero_indices)), k=batch_size)
             indices = nonzero_indices[chosen_indices]
         else:
-            indices = torch.floor(
-                torch.rand((batch_size, 3), device=device)
-                * torch.tensor([num_images, image_height, image_width], device=device)
-            ).long()
+            randids = torch.rand((batch_size, 3), device=device)* torch.tensor([num_images, image_height, image_width], device=device)
+            # indices = randids.floor()
+            # print("done floor")
+            indices = randids.long()
 
         return indices
 
@@ -101,13 +101,11 @@ class PixelSampler:
         collated_batch = {
             key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and value is not None
         }
-
         assert collated_batch["image"].shape[0] == num_rays_per_batch
 
         # Needed to correct the random indices to their actual camera idx locations.
         indices[:, 0] = batch["image_idx"][c]
         collated_batch["indices"] = indices  # with the abs camera indices
-
         if keep_full_image:
             collated_batch["full_image"] = batch["image"]
 

nerfstudio/fields/nerfacto_field.py

@@ -89,6 +89,7 @@ class NerfactoField(Field):
         use_transient_embedding: bool = False,
         use_semantics: bool = False,
         num_semantic_classes: int = 100,
+        freq_warmup: int = 0,
         pass_semantic_gradients: bool = False,
         use_pred_normals: bool = False,
         use_average_appearance_embedding: bool = False,
@@ -114,6 +115,8 @@ class NerfactoField(Field):
         self.use_pred_normals = use_pred_normals
         self.pass_semantic_gradients = pass_semantic_gradients
         self.base_res = base_res
+        self.freq_warmup = freq_warmup
+        self.step = 0
 
         self.direction_encoding = SHEncoding(
             levels=4,
@@ -198,6 +201,29 @@ class NerfactoField(Field):
             implementation=implementation,
         )
 
+    def step_cb(self, step):
+        self.step = step
+
+    def get_freq_mask(self, B, N, D, device):
+        import numpy as np
+
+        ts = np.linspace(self.step * N, self.step, N)
+        mask_vals = np.interp(ts, [0, self.freq_warmup], [0, 1])
+        m = torch.from_numpy(mask_vals)[None, ..., None].to(device)
+        return m.repeat((B, 1, D))
+
+    def mask_freqs(self, hashgrid_outputs):
+        """
+        hashgrid_outputs is a B x N*D array of the hashgrid values
+        """
+        # N is number of levels (usually 16)
+        # D is feature dim (usually 2)
+        N = self.mlp_base_grid.num_levels
+        D = self.mlp_base_grid.features_per_level
+        hashgrid_outputs = hashgrid_outputs.view(-1, N, D)
+        masked = self.get_freq_mask(hashgrid_outputs.shape[0], N, D, hashgrid_outputs.device) * hashgrid_outputs
+        return masked.view(masked.shape[0], -1)
+
     def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]:
         """Computes and returns the densities."""
         if self.spatial_distortion is not None:
@@ -213,7 +239,10 @@ class NerfactoField(Field):
         if not self._sample_locations.requires_grad:
             self._sample_locations.requires_grad = True
         positions_flat = positions.view(-1, 3)
-        h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1)
+        hashgrid_vecs = self.mlp_base_grid(positions_flat)
+        hashgrid_vecs = self.mask_freqs(hashgrid_vecs)
+        h = self.mlp_base_mlp(hashgrid_vecs).view(*ray_samples.frustums.shape, -1)
+        # h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1)
         density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1)
         self._density_before_activation = density_before_activation
 

nerfstudio/models/nerfacto.py

@@ -18,6 +18,7 @@ NeRF implementation that combines many recent advancements.
 
 from __future__ import annotations
 
+from contextlib import nullcontext
 from dataclasses import dataclass, field
 from typing import Dict, List, Literal, Tuple, Type
 
@@ -283,13 +284,19 @@ class NerfactoModel(Model):
                     func=self.proposal_sampler.step_cb,
                 )
             )
+            callbacks.append(
+                TrainingCallback(
+                    where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
+                    update_every_num_iters=1,
+                    func=self.field.step_cb,
+                )
+            )
         return callbacks
 
     def get_outputs(self, ray_bundle: RayBundle):
         # apply the camera optimizer pose tweaks
         if self.training:
             self.camera_optimizer.apply_to_raybundle(ray_bundle)
-
         ray_samples: RaySamples
         ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns)
         field_outputs = self.field.forward(ray_samples, compute_normals=self.config.predict_normals)
@@ -333,7 +340,6 @@ class NerfactoModel(Model):
 
         for i in range(self.config.num_proposal_iterations):
             outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i])
-
         return outputs
 
     def get_metrics_dict(self, outputs, batch):

nerfstudio/utils/tensor_dataclass.py

@@ -260,6 +260,14 @@ class TensorDataclass:
             A new TensorDataclass with the same data but on the specified device.
         """
         return self._apply_fn_to_fields(lambda x: x.to(device))
+    
+    def pin_memory(self: TensorDataclassT) -> TensorDataclassT:
+        """Pins the tensor dataclass memory
+
+        Returns:
+            TensorDataclass: A new TensorDataclass with the same data but pinned.
+        """
+        return self._apply_fn_to_fields(lambda x: x.pin_memory())
 
     def _apply_fn_to_fields(
         self: TensorDataclassT,

nerfstudio/viewer_beta/render_state_machine.py

@@ -181,8 +181,8 @@ class RenderStateMachine(threading.Thread):
                 # if we got interrupted, don't send the output to the viewer
                 continue
             self._send_output_to_viewer(outputs)
-            if self.state == "high":
-                self.viewer.update_camera_poses()
+            # if self.state == "high":
+            #     self.viewer.update_camera_poses()
             # if we rendered a static low res, we need to self-trigger a static high-res
             if self.next_action is None:
                 # if there hasn't been an action, wait for 1/target_fps seconds in case we get another move command

nerfstudio/viewer_beta/viewer.py

@@ -212,7 +212,7 @@ class Viewer:
             R = vtf.SO3.from_matrix(c2w[:3, :3])
             R = R @ vtf.SO3.from_x_radians(np.pi)
             camera_handle = self.viser_server.add_camera_frustum(
-                name=f"/camera_{idx:05d}",
+                name=f"/Train Cameras/camera_{idx:05d}",
                 fov=2 * np.arctan(camera.cx / camera.fx[0]),
                 scale=1,
                 aspect=float(camera.cx[0] / camera.cy[0]),