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

Option to process images on GPU for speed at the expense of memory (#2288)

Browse Source 
* Fix pixel sampler cpu bottleneck

cpu usage reduced 1700% -> 100%, and even more if using multi-gpu training

* Convert to config option

* Change optional to bool
This commit is contained in:
AX-I 2023-08-01 14:12:30 -04:00 committed by GitHub
parent 7e8d67f8a8
commit 98d36ed7d9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -115,8 +115,10 @@ class DataManagerConfig(InstantiateConfig):
"""Source of data, may not be used by all models."""
camera_optimizer: Optional[CameraOptimizerConfig] = None
"""Specifies the camera pose optimizer used during training. Helpful if poses are noisy."""
masks_on_gpu: Optional[bool] = None
masks_on_gpu: bool = False
"""Process masks on GPU for speed at the expense of memory, if True."""
images_on_gpu: bool = False
"""Process images on GPU for speed at the expense of memory, if True."""
class DataManager(nn.Module):
@ -403,6 +405,8 @@ class VanillaDataManager(DataManager, Generic[TDataset]):
self.exclude_batch_keys_from_device = self.train_dataset.exclude_batch_keys_from_device
if self.config.masks_on_gpu is True:
self.exclude_batch_keys_from_device.remove("mask")
if self.config.images_on_gpu is True:
self.exclude_batch_keys_from_device.remove("image")
if self.train_dataparser_outputs is not None:
cameras = self.train_dataparser_outputs.cameras