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

@@ -184,3 +184,4 @@ node_modules/
 bash/
 cache/
 renders/
+package-lock.json

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):

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

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}

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
 

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
 

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