Reduce GPU memory usage in CacheDataloader (#1730)

* Reduce GPU memory usage in CacheDataloader

* Move also masks and seg.masks, depths to cpu

* Fix linter issues

nerfstudio/data/datamanagers/base_datamanager.py

@@ -389,10 +389,11 @@ class VanillaDataManager(DataManager, Generic[TDataset]):
             self.config.data = self.config.dataparser.data
         self.dataparser = self.dataparser_config.setup()
         self.includes_time = self.dataparser.includes_time
-        self.train_dataparser_outputs = self.dataparser.get_dataparser_outputs(split="train")
+        self.train_dataparser_outputs: DataparserOutputs = self.dataparser.get_dataparser_outputs(split="train")
 
         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.train_dataparser_outputs is not None:
             cameras = self.train_dataparser_outputs.cameras
@@ -452,6 +453,7 @@ class VanillaDataManager(DataManager, Generic[TDataset]):
             num_workers=self.world_size * 4,
             pin_memory=True,
             collate_fn=self.config.collate_fn,
+            exclude_batch_keys_from_device=self.exclude_batch_keys_from_device,
         )
         self.iter_train_image_dataloader = iter(self.train_image_dataloader)
         self.train_pixel_sampler = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch)
@@ -475,6 +477,7 @@ class VanillaDataManager(DataManager, Generic[TDataset]):
             num_workers=self.world_size * 4,
             pin_memory=True,
             collate_fn=self.config.collate_fn,
+            exclude_batch_keys_from_device=self.exclude_batch_keys_from_device,
         )
         self.iter_eval_image_dataloader = iter(self.eval_image_dataloader)
         self.eval_pixel_sampler = self._get_pixel_sampler(self.eval_dataset, self.config.eval_num_rays_per_batch)

nerfstudio/data/datasets/base_dataset.py

@@ -42,6 +42,7 @@ class InputDataset(Dataset):
         scale_factor: The scaling factor for the dataparser outputs
     """
 
+    exclude_batch_keys_from_device: List[str] = ["image", "mask"]
     cameras: Cameras
 
     def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0):

nerfstudio/data/datasets/depth_dataset.py

@@ -31,6 +31,8 @@ class DepthDataset(InputDataset):
         scale_factor: The scaling factor for the dataparser outputs.
     """
 
+    exclude_batch_keys_from_device = InputDataset.exclude_batch_keys_from_device + ["depth_image"]
+
     def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0):
         super().__init__(dataparser_outputs, scale_factor)
         assert (

nerfstudio/data/datasets/sdf_dataset.py

@@ -34,6 +34,8 @@ class SDFDataset(InputDataset):
         scale_factor: The scaling factor for the dataparser outputs.
     """
 
+    exclude_batch_keys_from_device = InputDataset.exclude_batch_keys_from_device + ["depth", "normal"]
+
     def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0):
         super().__init__(dataparser_outputs, scale_factor)
 

nerfstudio/data/datasets/semantic_dataset.py

@@ -32,6 +32,8 @@ class SemanticDataset(InputDataset):
         dataparser_outputs: description of where and how to read input images.
     """
 
+    exclude_batch_keys_from_device = InputDataset.exclude_batch_keys_from_device + ["mask", "semantics"]
+
     def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0):
         super().__init__(dataparser_outputs, scale_factor)
         assert "semantics" in dataparser_outputs.metadata.keys() and isinstance(self.metadata["semantics"], Semantics)

nerfstudio/data/utils/dataloaders.py

@@ -21,7 +21,7 @@ import concurrent.futures
 import multiprocessing
 import random
 from abc import abstractmethod
-from typing import Any, Callable, Dict, Optional, Sized, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Sized, Tuple, Union
 
 import torch
 from rich.progress import track
@@ -55,8 +55,11 @@ class CacheDataloader(DataLoader):
         num_times_to_repeat_images: int = -1,
         device: Union[torch.device, str] = "cpu",
         collate_fn: Callable[[Any], Any] = nerfstudio_collate,
+        exclude_batch_keys_from_device: Optional[List[str]] = None,
         **kwargs,
     ):
+        if exclude_batch_keys_from_device is None:
+            exclude_batch_keys_from_device = ["image"]
         self.dataset = dataset
         assert isinstance(self.dataset, Sized)
 
@@ -67,6 +70,7 @@ class CacheDataloader(DataLoader):
         self.device = device
         self.collate_fn = collate_fn
         self.num_workers = kwargs.get("num_workers", 0)
+        self.exclude_batch_keys_from_device = exclude_batch_keys_from_device
 
         self.num_repeated = self.num_times_to_repeat_images  # starting value
         self.first_time = True
@@ -118,7 +122,9 @@ 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, exclude=["image"])
+        collated_batch = get_dict_to_torch(
+            collated_batch, device=self.device, exclude=self.exclude_batch_keys_from_device
+        )
         return collated_batch
 
     def __iter__(self):

nerfstudio/models/depth_nerfacto.py

@@ -107,7 +107,7 @@ class DepthNerfactoModel(NerfactoModel):
     ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
         """Appends ground truth depth to the depth image."""
         metrics, images = super().get_image_metrics_and_images(outputs, batch)
-        ground_truth_depth = batch["depth_image"]
+        ground_truth_depth = batch["depth_image"].to(self.device)
         if not self.config.is_euclidean_depth:
             ground_truth_depth = ground_truth_depth * outputs["directions_norm"]
 

nerfstudio/models/nerfplayer_nerfacto.py

@@ -248,11 +248,12 @@ class NerfplayerNerfactoModel(NerfactoModel):
                     outputs["rendered_pred_normal_loss"]
                 )
             if "depth_image" in batch.keys() and self.config.depth_weight > 0:
-                mask = (batch["depth_image"] != 0).view([-1])
+                depth_image = batch["depth_image"].to(self.device)
+                mask = (depth_image != 0).view([-1])
                 loss_dict["depth_loss"] = 0
 
                 def compute_depth_loss(x):
-                    return self.config.depth_weight * (x - batch["depth_image"][mask]).pow(2).mean()
+                    return self.config.depth_weight * (x - depth_image[mask]).pow(2).mean()
 
                 loss_dict["depth_loss"] = compute_depth_loss(outputs["depth"][mask])
                 for i in range(self.config.num_proposal_iterations):

nerfstudio/models/nerfplayer_ngp.py

@@ -226,9 +226,10 @@ class NerfplayerNGPModel(NGPModel):
         rgb_loss = self.rgb_loss(image, outputs["rgb"])
         loss_dict = {"rgb_loss": rgb_loss}
         if "depth_image" in batch.keys() and self.config.depth_weight > 0 and self.training:
-            mask = batch["depth_image"] != 0
+            depth_image = batch["depth_image"].to(self.device)
+            mask = depth_image != 0
             # First we calculate the depth value, just like most of the papers.
-            loss_dict["depth_loss"] = (outputs["depth"][mask] - batch["depth_image"][mask]).abs().mean()
+            loss_dict["depth_loss"] = (outputs["depth"][mask] - depth_image[mask]).abs().mean()
             # But this is not enough -- it will lead to fog like reconstructions, even with depth supervision.
             # Because this loss only cares about the mean value.
             # (Feel free to try it and see the fog-like effects on DyCheck by commenting out the following loss.)
@@ -244,7 +245,7 @@ class NerfplayerNGPModel(NGPModel):
             # "clear" loss (as it is directly applied to the network outputs, rather than post-processed values).
             # But such a loss also has drawbacks: it tends to overfit wrong (or noise) presented in the depth map.
             # Some structures are floating in the air with this loss...
-            gt_depth_packed = batch["depth_image"][outputs["ray_indices"]]
+            gt_depth_packed = depth_image[outputs["ray_indices"]]
             steps = (outputs["ray_samples"].frustums.starts + outputs["ray_samples"].frustums.ends) / 2
             # empty area should not be too close to the given depth, so lets add a margin to the gt depth
             margin = (self.scene_box.aabb.max() - self.scene_box.aabb.min()) / 128

nerfstudio/models/semantic_nerfw.py

@@ -254,7 +254,7 @@ class SemanticNerfWModel(Model):
 
         # semantic loss
         loss_dict["semantics_loss"] = self.config.semantic_loss_weight * self.cross_entropy_loss(
-            outputs["semantics"], batch["semantics"][..., 0].long()
+            outputs["semantics"], batch["semantics"][..., 0].long().to(self.device)
         )
         return loss_dict
 
@@ -301,6 +301,6 @@ class SemanticNerfWModel(Model):
         images_dict["semantics_colormap"] = self.colormap.to(self.device)[semantic_labels]
 
         # valid mask
-        images_dict["mask"] = batch["mask"].repeat(1, 1, 3)
+        images_dict["mask"] = batch["mask"].repeat(1, 1, 3).to(self.device)
 
         return metrics_dict, images_dict