From 3b0f7581d8df44045f0a0822f0cbcb7e6c0223b8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=C3=A1=C5=A1=20Kulh=C3=A1nek?= Date: Wed, 7 Jun 2023 20:14:24 +0200 Subject: [PATCH] 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 | 5 ++++- nerfstudio/data/datasets/base_dataset.py | 1 + nerfstudio/data/datasets/depth_dataset.py | 2 ++ nerfstudio/data/datasets/sdf_dataset.py | 2 ++ nerfstudio/data/datasets/semantic_dataset.py | 2 ++ nerfstudio/data/utils/dataloaders.py | 10 ++++++++-- nerfstudio/models/depth_nerfacto.py | 2 +- nerfstudio/models/nerfplayer_nerfacto.py | 5 +++-- nerfstudio/models/nerfplayer_ngp.py | 7 ++++--- nerfstudio/models/semantic_nerfw.py | 4 ++-- 10 files changed, 29 insertions(+), 11 deletions(-) diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index aeea43e3c..6b950de16 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/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) diff --git a/nerfstudio/data/datasets/base_dataset.py b/nerfstudio/data/datasets/base_dataset.py index ea7cfd1ef..95269a162 100644 --- a/nerfstudio/data/datasets/base_dataset.py +++ b/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): diff --git a/nerfstudio/data/datasets/depth_dataset.py b/nerfstudio/data/datasets/depth_dataset.py index 92adcfcc4..1d64cdf1c 100644 --- a/nerfstudio/data/datasets/depth_dataset.py +++ b/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 ( diff --git a/nerfstudio/data/datasets/sdf_dataset.py b/nerfstudio/data/datasets/sdf_dataset.py index 37d541220..a94ce04db 100644 --- a/nerfstudio/data/datasets/sdf_dataset.py +++ b/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) diff --git a/nerfstudio/data/datasets/semantic_dataset.py b/nerfstudio/data/datasets/semantic_dataset.py index 848dea8d0..ea209086f 100644 --- a/nerfstudio/data/datasets/semantic_dataset.py +++ b/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) diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index a6216d16d..e546a130b 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/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): diff --git a/nerfstudio/models/depth_nerfacto.py b/nerfstudio/models/depth_nerfacto.py index e2c7b6f05..e3328fd6d 100644 --- a/nerfstudio/models/depth_nerfacto.py +++ b/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"] diff --git a/nerfstudio/models/nerfplayer_nerfacto.py b/nerfstudio/models/nerfplayer_nerfacto.py index 46f0fcd3f..7a5a56c25 100644 --- a/nerfstudio/models/nerfplayer_nerfacto.py +++ b/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): diff --git a/nerfstudio/models/nerfplayer_ngp.py b/nerfstudio/models/nerfplayer_ngp.py index d820b0f54..c19cb4da2 100644 --- a/nerfstudio/models/nerfplayer_ngp.py +++ b/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 diff --git a/nerfstudio/models/semantic_nerfw.py b/nerfstudio/models/semantic_nerfw.py index c465e65b9..bf7914b1c 100644 --- a/nerfstudio/models/semantic_nerfw.py +++ b/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