tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 18:44:21 +08:00
Code

Fix dynamic batch logic during inference (#1160)

Browse Source
This commit is contained in:
Matthew Tancik 2022-12-21 22:50:05 -07:00 committed by GitHub
parent 8d4a82ef0b
commit 335c5c3661
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
nerfstudio/data/datamanagers/base_datamanager.py
View File

@ -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)

6
nerfstudio/pipelines/dynamic_batch.py
View File

@ -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,