From ab35e288711a5d9ec2fba3f8405d1b63f1b0fb58 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Fri, 15 Jul 2022 17:52:15 -0700 Subject: [PATCH] Cache full dataset by default (#157) Cache full dataset by default --- configs/graph_default.yaml | 3 +-- configs/graph_instant_ngp.yaml | 2 +- configs/graph_mipnerf.yaml | 9 --------- pyrad/data/image_sampler.py | 18 +++++++++++------- 4 files changed, 13 insertions(+), 19 deletions(-) diff --git a/configs/graph_default.yaml b/configs/graph_default.yaml index da824d4dd..176860e8d 100644 --- a/configs/graph_default.yaml +++ b/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 diff --git a/configs/graph_instant_ngp.yaml b/configs/graph_instant_ngp.yaml index c7a0821f0..935428eec 100644 --- a/configs/graph_instant_ngp.yaml +++ b/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 diff --git a/configs/graph_mipnerf.yaml b/configs/graph_mipnerf.yaml index 3f6963031..ac913bd2f 100644 --- a/configs/graph_mipnerf.yaml +++ b/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 diff --git a/pyrad/data/image_sampler.py b/pyrad/data/image_sampler.py index 96b464958..49b351513 100644 --- a/pyrad/data/image_sampler.py +++ b/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