From 50d9aa1341684fe6187e6c04bac73ea4f5d6b25c Mon Sep 17 00:00:00 2001 From: evonneng Date: Sun, 2 Oct 2022 13:54:13 -0700 Subject: [PATCH] small fixes (#581) * small fix * oom error * created model wrapper; make DDP working again; todo need to check if it works with viewer * working with multigpu now --- .gitignore | 1 + nerfstudio/data/utils/dataloaders.py | 6 +++-- nerfstudio/engine/trainer.py | 6 ++--- nerfstudio/pipelines/base.py | 31 ++++++++++++++---------- nerfstudio/utils/decorators.py | 4 +-- nerfstudio/utils/writer.py | 1 + nerfstudio/viewer/server/viewer_utils.py | 12 ++++++--- 7 files changed, 38 insertions(+), 23 deletions(-) diff --git a/.gitignore b/.gitignore index 51c91f933..6b2a94eba 100644 --- a/.gitignore +++ b/.gitignore @@ -184,3 +184,4 @@ node_modules/ bash/ cache/ renders/ +package-lock.json diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index a9bac84dd..0032b8d74 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/nerfstudio/data/utils/dataloaders.py @@ -21,6 +21,7 @@ from abc import abstractmethod from typing import Dict, Optional, Tuple, Union import torch +from rich.progress import track from torch.utils.data import Dataset, default_collate from torch.utils.data.dataloader import DataLoader @@ -69,8 +70,9 @@ class CacheDataloader(DataLoader): def _get_batch_list(self): """Returns a list of batches from the dataset attribute.""" indices = random.sample(range(len(self.dataset)), k=self.num_images_to_sample_from) - - batch_list = [self.dataset.__getitem__(idx) for idx in indices] + batch_list = [] + for idx in track(indices, description="Loading data batch"): + batch_list.append(self.dataset.__getitem__(idx)) return batch_list def _get_collated_batch(self): diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index f77087d21..0c1992fdd 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -106,10 +106,10 @@ class Trainer: self.prev_ckpt_paths = [] # set up viewer if enabled viewer_log_path = self.base_dir / config.viewer.relative_log_filename + ret = None if self.config.is_viewer_enabled(): - self.viewer_state, banner_messages = viewer_utils.setup_viewer(config.viewer, log_filename=viewer_log_path) - else: - self.viewer_state, banner_messages = None, None + ret = viewer_utils.setup_viewer(config.viewer, log_filename=viewer_log_path) + (self.viewer_state, banner_messages) = ret if ret else (None, None) self._check_viewer_warnings() # set up writers/profilers if enabled writer_log_path = self.base_dir / config.logging.relative_log_dir diff --git a/nerfstudio/pipelines/base.py b/nerfstudio/pipelines/base.py index 784cd2183..d962f3f08 100644 --- a/nerfstudio/pipelines/base.py +++ b/nerfstudio/pipelines/base.py @@ -90,12 +90,17 @@ class Pipeline(nn.Module): # pylint: disable=abstract-method datamanager: DataManager - model: Model + _model: Model + + @property + def model(self): + """Returns the unwrapped model if in ddp""" + return module_wrapper(self._model) @property def device(self): """Returns the device that the model is on.""" - return module_wrapper(self.model).device + return self.model.device @profiler.time_function def get_train_loss_dict(self, step: int): @@ -210,9 +215,9 @@ class VanillaPipeline(Pipeline): device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank ) self.datamanager.to(device) - # TODO(ethan): get rid of scene_box from the model + # TODO(ethan): get rid of scene_bounds from the model assert self.datamanager.train_dataset is not None, "Missing input dataset" - self.model = config.model.setup( + self._model = config.model.setup( scene_box=self.datamanager.train_dataset.dataset_inputs.scene_box, num_train_data=len(self.datamanager.train_dataset), ) @@ -220,15 +225,15 @@ class VanillaPipeline(Pipeline): self.world_size = world_size if world_size > 1: - self.model = typing.cast( - Model, typing.cast(Model, DDP(self.model, device_ids=[local_rank], find_unused_parameters=True)) + self._model = typing.cast( + Model, typing.cast(Model, DDP(self._model, device_ids=[local_rank], find_unused_parameters=True)) ) dist.barrier(device_ids=[local_rank]) @property def device(self): """Returns the device that the model is on.""" - return module_wrapper(self.model).device + return self.model.device @profiler.time_function def get_train_loss_dict(self, step: int): @@ -285,8 +290,8 @@ class VanillaPipeline(Pipeline): """ self.eval() image_idx, camera_ray_bundle, batch = self.datamanager.next_eval_image(step) - outputs = module_wrapper(self.model).get_outputs_for_camera_ray_bundle(camera_ray_bundle) - metrics_dict, images_dict = module_wrapper(self.model).get_image_metrics_and_images(outputs, batch) + outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + metrics_dict, images_dict = self.model.get_image_metrics_and_images(outputs, batch) assert "image_idx" not in metrics_dict metrics_dict["image_idx"] = image_idx assert "num_rays" not in metrics_dict @@ -317,8 +322,8 @@ class VanillaPipeline(Pipeline): inner_start = time() height, width = camera_ray_bundle.shape num_rays = height * width - outputs = module_wrapper(self.model).get_outputs_for_camera_ray_bundle(camera_ray_bundle) - metrics_dict, _ = module_wrapper(self.model).get_image_metrics_and_images(outputs, batch) + outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + metrics_dict, _ = self.model.get_image_metrics_and_images(outputs, batch) assert "num_rays_per_sec" not in metrics_dict metrics_dict["num_rays_per_sec"] = num_rays / (time() - inner_start) fps_str = f"fps_at_{height}x{width}" @@ -349,7 +354,7 @@ class VanillaPipeline(Pipeline): ) -> List[TrainingCallback]: """Returns the training callbacks from both the Dataloader and the Model.""" datamanager_callbacks = self.datamanager.get_training_callbacks(training_callback_attributes) - model_callbacks = module_wrapper(self.model).get_training_callbacks(training_callback_attributes) + model_callbacks = self.model.get_training_callbacks(training_callback_attributes) callbacks = datamanager_callbacks + model_callbacks return callbacks @@ -360,6 +365,6 @@ class VanillaPipeline(Pipeline): A list of dictionaries containing the pipeline's param groups. """ datamanager_params = self.datamanager.get_param_groups() - model_params = module_wrapper(self.model).get_param_groups() + model_params = self.model.get_param_groups() # TODO(ethan): assert that key names don't overlap return {**datamanager_params, **model_params} diff --git a/nerfstudio/utils/decorators.py b/nerfstudio/utils/decorators.py index 9e9cc51bf..91ffc19fb 100644 --- a/nerfstudio/utils/decorators.py +++ b/nerfstudio/utils/decorators.py @@ -50,11 +50,11 @@ def check_profiler_enabled(func: Callable) -> Callable: def check_viewer_enabled(func: Callable) -> Callable: - """Decorator: check if viewer is enabled""" + """Decorator: check if viewer is enabled and only run on main process""" def wrapper(self, *args, **kwargs): ret = None - if self.config.is_viewer_enabled(): + if self.config.is_viewer_enabled() and comms.is_main_process(): ret = func(self, *args, **kwargs) return ret diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index 5b6005f91..d0cf71713 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -190,6 +190,7 @@ def setup_local_writer(config: cfg.LoggingConfig, max_iter: int, banner_messages GLOBAL_BUFFER["events"] = {} +@check_main_thread def setup_event_writer(config: cfg.Config, log_dir: Path) -> None: """Initialization of all event writers specified in config diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index ebfd6506c..dcea0e037 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -34,6 +34,7 @@ from nerfstudio.configs import base_config as cfg from nerfstudio.data.utils.datasets import InputDataset from nerfstudio.models.base import Model from nerfstudio.utils import profiler, visualization, writer +from nerfstudio.utils.decorators import check_main_thread, decorate_all from nerfstudio.utils.io import load_from_json, write_to_json from nerfstudio.utils.misc import get_dict_to_torch from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter @@ -44,6 +45,7 @@ from nerfstudio.viewer.server.visualizer import Viewer console = Console(width=120) +@check_main_thread def setup_viewer(config: cfg.ViewerConfig, log_filename: Path): """Sets up the viewer if enabled @@ -193,6 +195,7 @@ class CheckThread(threading.Thread): return +@decorate_all([check_main_thread]) class ViewerState: """Class to hold state for viewer variables @@ -626,6 +629,7 @@ class ViewerState: check_thread = CheckThread(state=self) render_thread = RenderThread(state=self, graph=graph, camera_ray_bundle=camera_ray_bundle) + oom = False with TimeWriter(None, None, write=False) as vis_t: check_thread.start() render_thread.start() @@ -635,13 +639,15 @@ class ViewerState: except IOChangeException: pass except RuntimeError as e: - del camera_ray_bundle - torch.cuda.empty_cache() - time.sleep(0.5) # sleep to allow buffer to reset + oom = True self.vis["renderingState/log_errors"].write( "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing." ) print(f"Error: {e}") + if oom: + del camera_ray_bundle + torch.cuda.empty_cache() + time.sleep(0.5) # sleep to allow buffer to reset graph.train() outputs = render_thread.vis_outputs