improve memory consumption by clearing cache + move image to cpu (#412)

* reduced mem on train; figuring out memory leak on visualizer

* steadily below 10G memory

* cpu loading image

* update docstrings

* added scrolling

* Revert "update docstrings"

This reverts commit a914387cf8.

* Revert "Revert "update docstrings""

This reverts commit 177c3dd6f1.

* Revert "added scrolling"

This reverts commit f37351ee7f.

nerfactory/datamanagers/base.py

@@ -251,6 +251,8 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
                 self.train_input_dataset,
                 num_images_to_sample_from=self.config.train_num_images_to_sample_from,
                 device=self.device,
+                num_workers=self.world_size * 4,
+                pin_memory=True,
             )
         self.iter_train_image_dataloader = iter(self.train_image_dataloader)
         self.train_pixel_sampler = PixelSampler(self.config.train_num_rays_per_batch)

nerfactory/datamanagers/dataloaders.py

@@ -77,7 +77,7 @@ class CacheImageDataloader(DataLoader):
         """Returns a collated batch."""
         batch_list = self._get_batch_list()
         collated_batch = default_collate(batch_list)
-        collated_batch = get_dict_to_torch(collated_batch, device=self.device)
+        collated_batch = get_dict_to_torch(collated_batch, device=self.device, exclude=["image"])
         return collated_batch
 
     def __iter__(self):
@@ -162,7 +162,7 @@ class EvalDataloader(DataLoader):
         ray_bundle.num_rays_per_chunk = self.num_rays_per_chunk
         ray_bundle.camera_indices = torch.Tensor([image_idx])[..., None].int()
         batch = self.input_dataset[image_idx]
-        batch = get_dict_to_torch(batch, device=self.device)
+        batch = get_dict_to_torch(batch, device=self.device, exclude=["image"])
         return ray_bundle, batch
 
 

nerfactory/models/instant_ngp.py

@@ -150,7 +150,7 @@ class NGPModel(Model):
         return loss_dict
 
     def log_test_image_outputs(self, image_idx, step, batch, outputs):
-        image = batch["image"]
+        image = batch["image"].to(outputs["rgb"].device)
         rgb = outputs["rgb"]
         acc = visualization.apply_colormap(outputs["accumulation"])
         depth = visualization.apply_depth_colormap(

nerfactory/models/mipnerf.py

@@ -146,7 +146,7 @@ class MipNerfModel(Model):
         return loss_dict
 
     def log_test_image_outputs(self, image_idx, step, batch, outputs):
-        image = batch["image"]
+        image = batch["image"].to(outputs["rgb_coarse"].device)
         rgb_coarse = outputs["rgb_coarse"]
         rgb_fine = outputs["rgb_fine"]
         acc_coarse = visualization.apply_colormap(outputs["accumulation_coarse"])

nerfactory/models/mipnerf_360.py

@@ -161,7 +161,7 @@ class MipNerf360Model(Model):
         return loss_dict
 
     def log_test_image_outputs(self, image_idx, step, batch, outputs):
-        image = batch["image"]
+        image = batch["image"].to(outputs["rgb_coarse"].device)
         rgb_coarse = outputs["rgb_coarse"]
         rgb_fine = outputs["rgb_fine"]
         acc_coarse = visualization.apply_colormap(outputs["accumulation_coarse"])

nerfactory/models/nerfw.py

@@ -215,7 +215,7 @@ class NerfWModel(Model):
         return loss_dict
 
     def log_test_image_outputs(self, image_idx, step, batch, outputs):
-        image = batch["image"]
+        image = batch["image"].to(outputs["rgb_coarse"].device)
         rgb_coarse = outputs["rgb_coarse"]
         rgb_fine = outputs["rgb_fine"]
         rgb_fine_static = outputs["rgb_fine_static"]

nerfactory/models/semantic_nerf.py

@@ -114,7 +114,7 @@ class SemanticNerfModel(NeRFModel):
         return outputs
 
     def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients):
-        image = batch["image"]
+        image = batch["image"].to(outputs["rgb_coarse"].device)
         rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"])
         rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"])
         semantic_logits = outputs["semantic_fine"]

nerfactory/models/tensorf.py

@@ -196,7 +196,7 @@ class TensoRFModel(Model):
         return loss_dict
 
     def log_test_image_outputs(self, image_idx, step, batch, outputs):
-        image = batch["image"]
+        image = batch["image"].to(outputs["rgb"].device)
         rgb = outputs["rgb"]
         acc = visualization.apply_colormap(outputs["accumulation"])
         depth = visualization.apply_depth_colormap(

nerfactory/models/vanilla_nerf.py

@@ -179,7 +179,7 @@ class NeRFModel(Model):
         return loss_dict
 
     def log_test_image_outputs(self, image_idx, step, batch, outputs):
-        image = batch["image"]
+        image = batch["image"].to(outputs["rgb_coarse"].device)
         rgb_coarse = outputs["rgb_coarse"]
         rgb_fine = outputs["rgb_fine"]
         acc_coarse = visualization.apply_colormap(outputs["accumulation_coarse"])

nerfactory/utils/misc.py

@@ -19,20 +19,25 @@ Miscellaneous helper code.
 import hashlib
 import json
 from math import floor, log
-from typing import Any, Callable, Dict, Union
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import torch
 
 
-def get_dict_to_torch(stuff: Any, device: Union[torch.device, str] = "cpu"):
+def get_dict_to_torch(stuff: Any, device: Union[torch.device, str] = "cpu", exclude: Optional[List[str]] = None):
     """Set everything in the dict to the specified torch device.
 
     Args:
         stuff: things to convert to torch
+        device: machine to put the "stuff" on
+        exclude: list of keys to skip over transferring to device
     """
     if isinstance(stuff, dict):
         for k, v in stuff.items():
-            stuff[k] = get_dict_to_torch(v, device)
+            if exclude and k in exclude:
+                stuff[k] = v
+            else:
+                stuff[k] = get_dict_to_torch(v, device)
         return stuff
     if isinstance(stuff, torch.Tensor):
         return stuff.to(device)

nerfactory/viewer/server/viewer_utils.py

@@ -32,6 +32,7 @@ from nerfactory.datamanagers.datasets import InputDataset
 from nerfactory.models.base import Model
 from nerfactory.utils import profiler, visualization, writer
 from nerfactory.utils.decorators import check_visualizer_enabled, decorate_all
+from nerfactory.utils.misc import get_dict_to_torch
 from nerfactory.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter
 from nerfactory.viewer.server.subprocess import run_viewer_bridge_server_as_subprocess
 from nerfactory.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_h
@@ -107,6 +108,7 @@ class RenderThread(threading.Thread):
             self.exc = e
 
         if outputs:
+            outputs = get_dict_to_torch(outputs)
             self.vis_outputs = outputs
 
         self.state.check_done_render = True
@@ -114,6 +116,7 @@ class RenderThread(threading.Thread):
 
     def join(self, timeout=None):
         threading.Thread.join(self)
+        torch.cuda.empty_cache()
         if self.exc:
             raise self.exc