From 267924817793e0b0760c9eb2fb9ea9c8f0197fcb Mon Sep 17 00:00:00 2001
From: Justin Kerr <justin.g.kerr@gmail.com>
Date: Mon, 7 Nov 2022 23:19:21 -0800
Subject: [PATCH] refactor dataset...still slow

---
 nerfstudio/configs/method_configs.py          |   2 +-
 .../data/datamanagers/base_datamanager.py     | 189 ++++++++++--------
 nerfstudio/data/datasets/base_dataset.py      |  27 ++-
 nerfstudio/data/utils/dataloaders.py          |  91 +--------
 nerfstudio/model_components/ray_generators.py |   1 -
 nerfstudio/utils/tensor_dataclass.py          |  11 +-
 6 files changed, 147 insertions(+), 174 deletions(-)

diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py
index 1765f7c94..377c782aa 100644
--- a/nerfstudio/configs/method_configs.py
+++ b/nerfstudio/configs/method_configs.py
@@ -63,7 +63,7 @@ method_configs["nerfacto"] = Config(
     pipeline=VanillaPipelineConfig(
         datamanager=VanillaDataManagerConfig(
             dataparser=NerfstudioDataParserConfig(),
-            train_num_rays_per_batch=8000,
+            train_num_rays_per_batch=4096,
             eval_num_rays_per_batch=4096,
             camera_optimizer=CameraOptimizerConfig(
                 mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py
index b4f869b42..9c9c97baf 100644
--- a/nerfstudio/data/datamanagers/base_datamanager.py
+++ b/nerfstudio/data/datamanagers/base_datamanager.py
@@ -18,6 +18,7 @@ Data loader.
 
 from __future__ import annotations
 
+import functools
 from abc import abstractmethod
 from dataclasses import dataclass, field
 from typing import Dict, List, Optional, Tuple, Type, Union
@@ -28,6 +29,7 @@ from rich.progress import Console
 from torch import nn
 from torch.nn import Parameter
 from torch.utils.data import Dataset
+from torch.utils.data.dataloader import DataLoader
 from torch.utils.data.distributed import DistributedSampler
 
 from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
@@ -44,9 +46,9 @@ from nerfstudio.data.dataparsers.record3d_dataparser import Record3DDataParserCo
 from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.data.pixel_samplers import PixelSampler
 from nerfstudio.data.utils.dataloaders import (
-    CacheDataloader,
     FixedIndicesEvalDataloader,
     RandIndicesEvalDataloader,
+    _collate_images,
 )
 from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes
 from nerfstudio.model_components.ray_generators import RayGenerator
@@ -264,42 +266,51 @@ class VanillaDataManagerConfig(InstantiateConfig):
     """Specifies the camera pose optimizer used during training. Helpful if poses are noisy, such as for data from
     Record3D."""
 
-import torch.multiprocessing as mp
-import queue
+
+# import copy
+# import queue
 import time
-MP_CON = mp.get_context('forkserver')
-class Runner(MP_CON.Process):
-    def __init__(self, bundle_queue, train_dataset, device, config):
-        super().__init__()
-        self.train_dataset = train_dataset
-        self.device = device
-        self.bundle_queue = bundle_queue
-        self.config = config
-        self.daemon=True
-        self.start()
-        
-    def run(self):
-        self.train_image_dataloader = CacheDataloader(
-            self.train_dataset,
-            num_images_to_sample_from=self.config.train_num_images_to_sample_from,
-            num_times_to_repeat_images=self.config.train_num_times_to_repeat_images,
-            device=self.device,
-            num_workers=2,
-            pin_memory=False,
-        )
-        self.iter_train_image_dataloader = iter(self.train_image_dataloader)
-        self.train_pixel_sampler = PixelSampler(self.config.train_num_rays_per_batch)
-        self.train_ray_generator = RayGenerator(
-            self.train_dataset.dataparser_outputs.cameras.to(self.device)
-        ).to(self.device)
-        while True:
-            image_batch = next(self.iter_train_image_dataloader)
-            batch = self.train_pixel_sampler.sample(image_batch)
-            ray_indices = batch["indices"]
-            ray_bundle = self.train_ray_generator(ray_indices)
-            self.bundle_queue.put((ray_bundle.origins,ray_bundle.directions,ray_bundle.pixel_area,ray_bundle.camera_indices,batch))
-                
-        
+
+# import torch.multiprocessing as mp
+
+# MP_CON = mp.get_context("forkserver")
+
+
+# class Runner(MP_CON.Process):
+#     def __init__(self, bundle_queue, train_dataset, device, config):
+#         super().__init__()
+#         self.train_dataset = train_dataset
+#         self.device = device
+#         self.bundle_queue = bundle_queue
+#         self.config = config
+#         self.daemon = True
+#         self.start()
+
+#     def run(self):
+#         self.train_image_dataloader = CacheDataloader(
+#             self.train_dataset,
+#             num_images_to_sample_from=self.config.train_num_images_to_sample_from,
+#             num_times_to_repeat_images=self.config.train_num_times_to_repeat_images,
+#             device=self.device,
+#             num_workers=1,
+#             pin_memory=False,
+#         )
+#         while True:
+#             image_batch = next(self.iter_train_image_dataloader)
+#             batch = self.train_pixel_sampler.sample(image_batch)
+#             ray_indices = batch["indices"]
+#             ray_bundle = self.train_ray_generator(ray_indices)
+#             self.bundle_queue.put(
+#                 (ray_bundle.origins, ray_bundle.directions, ray_bundle.pixel_area, ray_bundle.camera_indices, batch)
+#             )
+#             time.sleep(10)
+#             exit()
+
+
+def _test_collate(batch):
+    return batch[0]
+
+
 class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
     """Basic stored data manager implementation.
 
@@ -335,8 +346,8 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
 
         self.train_dataset = self.create_train_dataset()
         self.eval_dataset = self.create_eval_dataset()
-        self.bundle_queue = MP_CON.SimpleQueue()
-        runners = [Runner(self.bundle_queue, self.train_dataset, self.device,self.config) for i in range(2)]
+        # self.bundle_queue = MP_CON.SimpleQueue()
+        # runners = [Runner(self.bundle_queue, self.train_dataset, self.device, self.config) for i in range(1)]
         DataManager.__init__(self)
 
     def create_train_dataset(self) -> InputDataset:
@@ -348,7 +359,7 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         return InputDataset(
             self.config.dataparser.setup().get_dataparser_outputs(split="val" if not self.test_mode else "test")
         )
-        
+
     def setup_train(self):
         """Sets up the data loaders for training"""
         assert self.train_dataset is not None
@@ -356,71 +367,83 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         self.train_camera_optimizer = self.config.camera_optimizer.setup(
             num_cameras=self.train_dataset.dataparser_outputs.cameras.size, device=self.device
         )
+        self.train_image_dataloader = DataLoader(
+            self.train_dataset,
+            num_workers=4,
+            pin_memory=True,
+            collate_fn=_test_collate,
+            batch_size=1,
+            persistent_workers=True,
+        )
+        self.iter_train_image_dataloader = iter(self.train_image_dataloader)
+        # self.dummy = ray_bundle, batch
 
     def setup_eval(self):
         """Sets up the data loader for evaluation"""
         assert self.eval_dataset is not None
         CONSOLE.print("Setting up evaluation dataset...")
-        self.eval_image_dataloader = CacheDataloader(
+        self.eval_image_dataloader = DataLoader(
             self.eval_dataset,
-            num_images_to_sample_from=self.config.eval_num_images_to_sample_from,
-            num_times_to_repeat_images=self.config.eval_num_times_to_repeat_images,
-            device=self.device,
-            num_workers=self.world_size * 4,
+            num_workers=4,
             pin_memory=True,
+            collate_fn=_test_collate,
+            batch_size=1,
+            persistent_workers=True,
         )
         self.iter_eval_image_dataloader = iter(self.eval_image_dataloader)
-        self.eval_pixel_sampler = PixelSampler(self.config.eval_num_rays_per_batch)
-        self.eval_ray_generator = RayGenerator(
-            self.eval_dataset.dataparser_outputs.cameras.to(self.device)
-        )
-        # for loading full images
-        self.fixed_indices_eval_dataloader = FixedIndicesEvalDataloader(
-            input_dataset=self.eval_dataset,
-            device=self.device,
-            num_workers=self.world_size * 4,
-        )
-        self.eval_dataloader = RandIndicesEvalDataloader(
-            input_dataset=self.eval_dataset,
-            image_indices=self.config.eval_image_indices,
-            device=self.device,
-            num_workers=self.world_size * 4,
-        )
 
     def next_train(self, step: int) -> Tuple[RayBundle, Dict]:
         """Returns the next batch of data from the train dataloader."""
         self.train_count += 1
-        os,ds,ps,cs,batch = self.bundle_queue.get()
-        ray_bundle = RayBundle(
-            origins=os, directions=ds, pixel_area=ps, camera_indices=cs
-        )
-        # ray_bundle=ray_bundle.to(self.device)
+
+        # ray_bundle, batch = self.dummy
+        try:
+            ray_bundle, batch = next(self.iter_train_image_dataloader)
+        except StopIteration:
+            self.iter_train_image_dataloader = iter(self.train_image_dataloader)
+            ray_bundle, batch = next(self.iter_train_image_dataloader)
+        ray_bundle = ray_bundle.to(self.device)
         ray_indices = ray_bundle.camera_indices
-        
+
         # Multiply in the camera optimization deltas to the rays
-        cam_deltas = self.train_camera_optimizer(ray_indices[:,0])# B x 3 x 4
-        origins_homogeneous = torch.cat([ray_bundle.origins,torch.ones((*ray_bundle.shape,1),
-                                        dtype=ray_bundle.origins.dtype,device=ray_bundle.origins.device)],dim=-1)
-        origins_homogeneous = origins_homogeneous.view(*origins_homogeneous.shape,1)
-        ray_bundle.origins = torch.bmm(cam_deltas,origins_homogeneous)[...,:3,0]
-        ray_bundle.directions = torch.bmm(cam_deltas[:,:,:3],ray_bundle.directions.view(*ray_bundle.directions.shape,1))[...,0]
+        cam_deltas = self.train_camera_optimizer(ray_indices[:, 0])  # B x 3 x 4
+        origins_homogeneous = torch.cat(
+            [
+                ray_bundle.origins,
+                torch.ones((*ray_bundle.shape, 1), dtype=ray_bundle.origins.dtype, device=ray_bundle.origins.device),
+            ],
+            dim=-1,
+        )
+        origins_homogeneous = origins_homogeneous.view(*origins_homogeneous.shape, 1)
+        ray_bundle.origins = torch.bmm(cam_deltas, origins_homogeneous)[..., :3, 0]
+        ray_bundle.directions = torch.bmm(
+            cam_deltas[:, :, :3], ray_bundle.directions.view(*ray_bundle.directions.shape, 1)
+        )[..., 0]
         return ray_bundle, batch
 
     def next_eval(self, step: int) -> Tuple[RayBundle, Dict]:
         """Returns the next batch of data from the eval dataloader."""
         self.eval_count += 1
-        image_batch = next(self.iter_eval_image_dataloader)
-        batch = self.eval_pixel_sampler.sample(image_batch)
-        ray_indices = batch["indices"]
-        ray_bundle = self.eval_ray_generator(ray_indices)
-        
+        ray_bundle, batch = next(self.iter_eval_image_dataloader)
+        # batch = self.eval_pixel_sampler.sample(image_batch)
+        # ray_indices = batch["indices"]
+        # ray_bundle = self.eval_ray_generator(ray_indices)
+
         # Multiply in the camera optimization deltas to the rays
-        cam_deltas = self.train_camera_optimizer(ray_indices[:,0])# B x 3 x 4
-        origins_homogeneous = torch.cat([ray_bundle.origins,torch.ones((*ray_bundle.shape,1),
-                                        dtype=ray_bundle.origins.dtype,device=ray_bundle.origins.device)],dim=-1)
-        origins_homogeneous = origins_homogeneous.view(*origins_homogeneous.shape,1)
-        ray_bundle.origins = torch.bmm(cam_deltas,origins_homogeneous)[...,:3,0]
-        ray_bundle.directions = torch.bmm(cam_deltas[:,:,:3],ray_bundle.directions.view(*ray_bundle.directions.shape,1))[...,0]
+        ray_indices = ray_bundle.camera_indices
+        cam_deltas = self.train_camera_optimizer(ray_indices[:, 0])  # B x 3 x 4
+        origins_homogeneous = torch.cat(
+            [
+                ray_bundle.origins,
+                torch.ones((*ray_bundle.shape, 1), dtype=ray_bundle.origins.dtype, device=ray_bundle.origins.device),
+            ],
+            dim=-1,
+        )
+        origins_homogeneous = origins_homogeneous.view(*origins_homogeneous.shape, 1)
+        ray_bundle.origins = torch.bmm(cam_deltas, origins_homogeneous)[..., :3, 0]
+        ray_bundle.directions = torch.bmm(
+            cam_deltas[:, :, :3], ray_bundle.directions.view(*ray_bundle.directions.shape, 1)
+        )[..., 0]
         return ray_bundle, batch
 
     def next_eval_image(self, step: int) -> Tuple[int, RayBundle, Dict]:
diff --git a/nerfstudio/data/datasets/base_dataset.py b/nerfstudio/data/datasets/base_dataset.py
index 0c6c9fe15..f10fb4995 100644
--- a/nerfstudio/data/datasets/base_dataset.py
+++ b/nerfstudio/data/datasets/base_dataset.py
@@ -17,20 +17,23 @@ Some dataset code.
 """
 from __future__ import annotations
 
+import random
 from typing import Dict
 
 import numpy as np
 import numpy.typing as npt
 import torch
 from PIL import Image
-from torch.utils.data import Dataset
+from torch.utils.data import IterableDataset, default_collate
 from torchtyping import TensorType
 
 from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs
+from nerfstudio.data.pixel_samplers import PixelSampler
 from nerfstudio.data.utils.data_utils import get_image_mask_tensor_from_path
+from nerfstudio.model_components.ray_generators import RayGenerator
 
 
-class InputDataset(Dataset):
+class InputDataset(IterableDataset):
     """Dataset that returns images.
 
     Args:
@@ -41,6 +44,15 @@ class InputDataset(Dataset):
         super().__init__()
         self.dataparser_outputs = dataparser_outputs
         self.has_masks = self.dataparser_outputs.mask_filenames is not None
+        self.cache = []
+        self.train_pixel_sampler = PixelSampler(4096)
+        self.train_ray_generator = RayGenerator(self.dataparser_outputs.cameras.to("cpu"))
+        print("Caching images...")
+        for i in range(len(self)):
+            data = self.get_data(i)
+            self.cache.append(data)
+        self.cache = default_collate(self.cache)
+        print("...done")
 
     def __len__(self):
         return len(self.dataparser_outputs.image_filenames)
@@ -99,6 +111,11 @@ class InputDataset(Dataset):
         del data
         return {}
 
-    def __getitem__(self, image_idx: int) -> Dict:
-        data = self.get_data(image_idx)
-        return data
+    def __getitem__(self, index):
+        return self.get_data(index)
+
+    def __iter__(self):
+        batch = self.train_pixel_sampler.sample(self.cache)
+        ray_indices = batch["indices"]
+        ray_bundle = self.train_ray_generator(ray_indices)
+        yield ray_bundle, batch
diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py
index 8f8873d65..d9aa70e03 100644
--- a/nerfstudio/data/utils/dataloaders.py
+++ b/nerfstudio/data/utils/dataloaders.py
@@ -28,93 +28,20 @@ from torch.utils.data.dataloader import DataLoader
 from nerfstudio.cameras.cameras import Cameras
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.data.datasets.base_dataset import InputDataset
+from nerfstudio.data.pixel_samplers import PixelSampler
+from nerfstudio.model_components.ray_generators import RayGenerator
 from nerfstudio.utils.misc import get_dict_to_torch
 
 CONSOLE = Console(width=120)
 
 
-class CacheDataloader(DataLoader):
-    """Collated image dataset that implements caching of default-pytorch-collatable data.
-    Creates batches of the InputDataset return type.
-
-    Args:
-        dataset: Dataset to sample from.
-        num_samples_to_collate: How many images to sample rays for each batch. -1 for all images.
-        num_times_to_repeat_images: How often to collate new images. -1 to never pick new images.
-        device: Device to perform computation.
-    """
-
-    def __init__(
-        self,
-        dataset: Dataset,
-        num_images_to_sample_from: int = -1,
-        num_times_to_repeat_images: int = -1,
-        device: Union[torch.device, str] = "cpu",
-        **kwargs,
-    ):
-        self.dataset = dataset
-        self.num_times_to_repeat_images = num_times_to_repeat_images
-        self.cache_all_images = (num_images_to_sample_from == -1) or (num_images_to_sample_from >= len(self.dataset))
-        self.num_images_to_sample_from = len(self.dataset) if self.cache_all_images else num_images_to_sample_from
-        self.device = device
-
-        self.num_repeated = self.num_times_to_repeat_images  # starting value
-        self.first_time = True
-
-        self.cached_collated_batch = None
-        if self.cache_all_images:
-            CONSOLE.print(f"Caching all {len(self.dataset)} images.")
-            if len(self.dataset) > 500:
-                CONSOLE.print(
-                    "[bold yellow]Warning: If you run out of memory, try reducing the number of images to sample from."
-                )
-            self.cached_collated_batch = self._get_collated_batch()
-        elif self.num_times_to_repeat_images == -1:
-            CONSOLE.print(
-                f"Caching {self.num_images_to_sample_from} out of {len(self.dataset)} images, without resampling."
-            )
-        else:
-            CONSOLE.print(
-                f"Caching {self.num_images_to_sample_from} out of {len(self.dataset)} images, "
-                f"resampling every {self.num_times_to_repeat_images} iters."
-            )
-        super().__init__(dataset=dataset, **kwargs)
-
-    def __getitem__(self, idx):
-        return self.dataset.__getitem__(idx)
-
-    def _get_batch_list(self):
-        """Returns a list of batches from the dataset attribute."""
-        indices = random.sample(range(len(self.dataset)), k=self.num_images_to_sample_from)
-        batch_list = []
-        for idx in track(indices, description="Loading data batch"):
-            batch_list.append(self.dataset.__getitem__(idx))
-        return batch_list
-
-    def _get_collated_batch(self):
-        """Returns a collated batch."""
-        batch_list = self._get_batch_list()
-        collated_batch = default_collate(batch_list)
-        collated_batch = get_dict_to_torch(collated_batch, device=self.device, exclude=["image"])
-        return collated_batch
-
-    def __iter__(self):
-        while True:
-            if self.cache_all_images:
-                collated_batch = self.cached_collated_batch
-            elif self.first_time or (
-                self.num_times_to_repeat_images != -1 and self.num_repeated >= self.num_times_to_repeat_images
-            ):
-                # trigger a reset
-                self.num_repeated = 0
-                collated_batch = self._get_collated_batch()
-                # possibly save a cached item
-                self.cached_collated_batch = collated_batch if self.num_times_to_repeat_images != 0 else None
-                self.first_time = False
-            else:
-                collated_batch = self.cached_collated_batch
-                self.num_repeated += 1
-            yield collated_batch
+def _collate_images(train_pixel_sampler, train_ray_generator, dataset_outputs):
+    with torch.no_grad():
+        batch = default_collate(dataset_outputs)
+        batch = train_pixel_sampler.sample(batch)
+        ray_indices = batch["indices"]
+        ray_bundle = train_ray_generator(ray_indices)
+        return ray_bundle.origins,ray_bundle.directions,ray_bundle.pixel_area,ray_bundle.camera_indices, batch
 
 
 class EvalDataloader(DataLoader):
diff --git a/nerfstudio/model_components/ray_generators.py b/nerfstudio/model_components/ray_generators.py
index b9658466a..03bc77fdd 100644
--- a/nerfstudio/model_components/ray_generators.py
+++ b/nerfstudio/model_components/ray_generators.py
@@ -49,7 +49,6 @@ class RayGenerator(nn.Module):
         x = ray_indices[:, 2]  # col indices
         coords = self.image_coords[y, x]
 
-
         ray_bundle = self.cameras.generate_rays(
             camera_indices=c,
             coords=coords,
diff --git a/nerfstudio/utils/tensor_dataclass.py b/nerfstudio/utils/tensor_dataclass.py
index cd85eac1c..20928c087 100644
--- a/nerfstudio/utils/tensor_dataclass.py
+++ b/nerfstudio/utils/tensor_dataclass.py
@@ -197,7 +197,7 @@ class TensorDataclass:
         """
         return self._apply_fn_to_fields(lambda x: x.broadcast_to((*shape, x.shape[-1])))
 
-    def to(self: TensorDataclassT, device) -> TensorDataclassT:
+    def to(self: TensorDataclassT, device, *args, **kwargs) -> TensorDataclassT:
         """Returns a new TensorDataclass with the same data but on the specified device.
 
         Args:
@@ -206,7 +206,14 @@ class TensorDataclass:
         Returns:
             A new TensorDataclass with the same data but on the specified device.
         """
-        return self._apply_fn_to_fields(lambda x: x.to(device))
+        return self._apply_fn_to_fields(lambda x: x.to(device, *args, **kwargs))
+
+    def detach(self: TensorDataclassT) -> TensorDataclassT:
+
+        return self._apply_fn_to_fields(lambda x: x.detach())
+
+    def pin_memory(self: TensorDataclassT) -> TensorDataclassT:
+        return self._apply_fn_to_fields(lambda x: x.pin_memory())
 
     def _apply_fn_to_fields(
         self: TensorDataclassT, fn: Callable, dataclass_fn: Optional[Callable] = None