tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

Avoid loading entire dataset on GPU (#1551)

Browse Source
This commit is contained in:
Matthew Tancik 2023-03-02 14:32:27 -08:00 committed by GitHub
parent a36f33c5ee
commit 8d41e522e3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
nerfstudio/data/utils/dataloaders.py
View File

@ -116,7 +116,7 @@ class CacheDataloader(DataLoader):
"""Returns a collated batch."""
batch_list = self._get_batch_list()
collated_batch = self.collate_fn(batch_list)
collated_batch = get_dict_to_torch(collated_batch, device=self.device)
collated_batch = get_dict_to_torch(collated_batch, device=self.device, exclude=["image"])
return collated_batch
def __iter__(self):
@ -183,7 +183,7 @@ class EvalDataloader(DataLoader):
"""
ray_bundle = self.cameras.generate_rays(camera_indices=image_idx, keep_shape=True)
batch = self.input_dataset[image_idx]
batch = get_dict_to_torch(batch, device=self.device)
batch = get_dict_to_torch(batch, device=self.device, exclude=["image"])
return ray_bundle, batch