From 335c5c366154edd4c47fc4899ee9b18c59e1ed10 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Wed, 21 Dec 2022 22:50:05 -0700 Subject: [PATCH] Fix dynamic batch logic during inference (#1160) --- nerfstudio/data/datamanagers/base_datamanager.py | 4 ++++ nerfstudio/pipelines/dynamic_batch.py | 6 ++++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 49b759a9d..26642a87c 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -298,6 +298,8 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method train_dataset: InputDataset eval_dataset: InputDataset train_dataparser_outputs: DataparserOutputs + train_pixel_sampler: Optional[PixelSampler] = None + eval_pixel_sampler: Optional[PixelSampler] = None def __init__( self, @@ -408,6 +410,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method """Returns the next batch of data from the train dataloader.""" self.train_count += 1 image_batch = next(self.iter_train_image_dataloader) + assert self.train_pixel_sampler is not None batch = self.train_pixel_sampler.sample(image_batch) ray_indices = batch["indices"] ray_bundle = self.train_ray_generator(ray_indices) @@ -417,6 +420,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method """Returns the next batch of data from the eval dataloader.""" self.eval_count += 1 image_batch = next(self.iter_eval_image_dataloader) + assert self.eval_pixel_sampler is not None batch = self.eval_pixel_sampler.sample(image_batch) ray_indices = batch["indices"] ray_bundle = self.eval_ray_generator(ray_indices) diff --git a/nerfstudio/pipelines/dynamic_batch.py b/nerfstudio/pipelines/dynamic_batch.py index 788f68e82..293ea1706 100644 --- a/nerfstudio/pipelines/dynamic_batch.py +++ b/nerfstudio/pipelines/dynamic_batch.py @@ -64,8 +64,10 @@ class DynamicBatchPipeline(VanillaPipeline): def _update_pixel_samplers(self): """Update the pixel samplers for train and eval with the dynamic number of rays per batch.""" - self.datamanager.train_pixel_sampler.set_num_rays_per_batch(self.dynamic_num_rays_per_batch) - self.datamanager.eval_pixel_sampler.set_num_rays_per_batch(self.dynamic_num_rays_per_batch) + if self.datamanager.train_pixel_sampler is not None: + self.datamanager.train_pixel_sampler.set_num_rays_per_batch(self.dynamic_num_rays_per_batch) + if self.datamanager.eval_pixel_sampler is not None: + self.datamanager.eval_pixel_sampler.set_num_rays_per_batch(self.dynamic_num_rays_per_batch) def _update_dynamic_num_rays_per_batch(self, num_samples_per_batch: int): """Updates the dynamic number of rays per batch variable,