Cache full dataset by default (#157)

Cache full dataset by default

configs/graph_default.yaml

@@ -58,8 +58,7 @@ data:
     _target_: pyrad.data.dataloader.TrainDataloader
     image_sampler:
       _target_: pyrad.data.image_sampler.CacheImageSampler
-      num_images_to_sample_from: 8
-      num_times_to_repeat_images: 100
+      num_images_to_sample_from: -1
     pixel_sampler:
       _target_: pyrad.data.pixel_sampler.PixelSampler
       num_rays_per_batch: 1024

configs/graph_instant_ngp.yaml

@@ -27,5 +27,5 @@ optimizers:
   fields:
     optimizer:
       lr: 1e-2
-      eps: 1e-4 # for float16 stability
+      eps: 1e-7 # need to investigate further for float16 stability
     scheduler: null

configs/graph_mipnerf.yaml

@@ -16,12 +16,3 @@ graph:
   loss_coefficients:
     rgb_loss_coarse: 0.1
     rgb_loss_fine: 1.0
-
-data:
-  # Train dataloading.
-  dataloader_train:
-    _target_: pyrad.data.dataloader.TrainDataloader
-    image_sampler:
-      _target_: pyrad.data.image_sampler.CacheImageSampler
-      num_images_to_sample_from: 0
-      num_times_to_repeat_images: -1

pyrad/data/image_sampler.py

@@ -44,7 +44,7 @@ class CacheImageSampler(ImageSampler):
 
     Args:
         dataset (torch.utils.data.Dataset): Dataset to sample from.
-        num_samples_to_collate (int): How many images are used to sample rays for each batch.
+        num_samples_to_collate (int, optional): How many images to sample rays for each batch. Defaults to all images.
         num_times_to_repeat_images (int): How often to collate new images. Defaults to every iteration.
         device (Union[torch.device, str]): Device to perform computation. Defaults to "cpu".
     """
@@ -52,22 +52,24 @@ class CacheImageSampler(ImageSampler):
     def __init__(
         self,
         dataset: torch.utils.data.Dataset,
-        num_images_to_sample_from: int,
+        num_images_to_sample_from: int = -1,
         num_times_to_repeat_images: int = 0,
         device: Union[torch.device, str] = "cpu",
     ):
         super().__init__()
         self.dataset = dataset
-        self.num_images_to_sample_from = (
-            len(self.dataset) if num_images_to_sample_from == 0 else num_images_to_sample_from
-        )
+        self.cache_all_images = num_images_to_sample_from is -1
+        self.num_images_to_sample_from = len(self.dataset) if self.cache_all_images else num_images_to_sample_from
         self.num_times_to_repeat_images = num_times_to_repeat_images
         self.device = device
 
         self.num_repeated = self.num_times_to_repeat_images  # starting value
-        self.cached_collated_batch = None
         self.first_time = True
 
+        self.cached_collated_batch = None
+        if self.cache_all_images:
+            self.cached_collated_batch = self._get_collated_batch()
+
     def __getitem__(self, idx):
         return self.dataset.__getitem__(idx)
 
@@ -86,7 +88,9 @@ class CacheImageSampler(ImageSampler):
 
     def __iter__(self):
         while True:
-            if self.first_time or (
+            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