From ccb8a25d4abdc4525052528904850b9a5b958a33 Mon Sep 17 00:00:00 2001 From: machenmusik Date: Wed, 21 Jun 2023 12:16:14 -0400 Subject: [PATCH] Add option to process masks on GPU for speed at the expense of memory (#2110) --pipeline.datamanager.masks-on-gpu True --- nerfstudio/data/datamanagers/base_datamanager.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 82e4f66d3..b2f23feb7 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -111,6 +111,8 @@ 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 + """Process masks on GPU for speed at the expense of memory, if True.""" class DataManager(nn.Module): @@ -396,6 +398,8 @@ class VanillaDataManager(DataManager, Generic[TDataset]): self.train_dataset = self.create_train_dataset() self.eval_dataset = self.create_eval_dataset() 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.train_dataparser_outputs is not None: cameras = self.train_dataparser_outputs.cameras