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

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