diff --git a/nerfactory/datamanagers/base.py b/nerfactory/datamanagers/base.py index ab7b3162a..fffa040d2 100644 --- a/nerfactory/datamanagers/base.py +++ b/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) diff --git a/nerfactory/datamanagers/dataloaders.py b/nerfactory/datamanagers/dataloaders.py index c58f7e400..8292828cd 100644 --- a/nerfactory/datamanagers/dataloaders.py +++ b/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 diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index eeffd5ac2..986e10a9d 100644 --- a/nerfactory/models/instant_ngp.py +++ b/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( diff --git a/nerfactory/models/mipnerf.py b/nerfactory/models/mipnerf.py index 514673c46..1e7772c76 100644 --- a/nerfactory/models/mipnerf.py +++ b/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"]) diff --git a/nerfactory/models/mipnerf_360.py b/nerfactory/models/mipnerf_360.py index 19dd35452..f9807efac 100644 --- a/nerfactory/models/mipnerf_360.py +++ b/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"]) diff --git a/nerfactory/models/nerfw.py b/nerfactory/models/nerfw.py index 176e2b001..d5e49fe05 100644 --- a/nerfactory/models/nerfw.py +++ b/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"] diff --git a/nerfactory/models/semantic_nerf.py b/nerfactory/models/semantic_nerf.py index b43d06d30..6909d9c7a 100644 --- a/nerfactory/models/semantic_nerf.py +++ b/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"] diff --git a/nerfactory/models/tensorf.py b/nerfactory/models/tensorf.py index 5817200cf..bac7089c2 100644 --- a/nerfactory/models/tensorf.py +++ b/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( diff --git a/nerfactory/models/vanilla_nerf.py b/nerfactory/models/vanilla_nerf.py index be2cdd431..f261274a3 100644 --- a/nerfactory/models/vanilla_nerf.py +++ b/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"]) diff --git a/nerfactory/utils/misc.py b/nerfactory/utils/misc.py index 0271e4d12..297c5d536 100644 --- a/nerfactory/utils/misc.py +++ b/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) diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index 952f1da84..8e23704c9 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/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