From bfaf11cd7259d48e877b373342c450702284d711 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Mon, 27 Feb 2023 18:32:40 -0800 Subject: [PATCH] Remove dataparser dependency (#1529) --- nerfstudio/configs/experiment_config.py | 11 ++-- .../data/datamanagers/base_datamanager.py | 55 ++++++++++++++++--- nerfstudio/engine/trainer.py | 50 ++++++++++------- nerfstudio/pipelines/base_pipeline.py | 22 ++++---- scripts/train.py | 4 +- scripts/viewer/run_viewer.py | 10 ++-- 6 files changed, 103 insertions(+), 49 deletions(-) diff --git a/nerfstudio/configs/experiment_config.py b/nerfstudio/configs/experiment_config.py index 98703ee40..72362befe 100644 --- a/nerfstudio/configs/experiment_config.py +++ b/nerfstudio/configs/experiment_config.py @@ -72,7 +72,7 @@ class ExperimentConfig(InstantiateConfig): vis: Literal["viewer", "wandb", "tensorboard", "viewer+wandb", "viewer+tensorboard"] = "wandb" """Which visualizer to use.""" data: Optional[Path] = None - """Alias for --pipeline.datamanager.dataparser.data""" + """Alias for --pipeline.datamanager.data""" relative_model_dir: Path = Path("nerfstudio_models/") """Relative path to save all checkpoints.""" @@ -96,9 +96,12 @@ class ExperimentConfig(InstantiateConfig): def set_experiment_name(self) -> None: """Dynamically set the experiment name""" if self.experiment_name is None: - datapath = self.pipeline.datamanager.dataparser.data - datapath = datapath.parent if datapath.is_file() else datapath - self.experiment_name = str(datapath).replace("../", "").replace("/", "-") + datapath = self.pipeline.datamanager.data + if datapath is not None: + datapath = datapath.parent if datapath.is_file() else datapath + self.experiment_name = str(datapath.stem) + else: + self.experiment_name = "unnamed" def get_base_dir(self) -> Path: """Retrieve the base directory to set relative paths""" diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 72137fb8e..dd0109c1d 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -20,6 +20,7 @@ from __future__ import annotations from abc import abstractmethod from dataclasses import dataclass, field +from pathlib import Path from typing import Any, Dict, List, Optional, Tuple, Type, Union import torch @@ -95,6 +96,21 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of same as the vanilla union, but results in shorter subcommand names.""" +@dataclass +class DataManagerConfig(InstantiateConfig): + """Configuration for data manager instantiation; DataManager is in charge of keeping the train/eval dataparsers; + After instantiation, data manager holds both train/eval datasets and is in charge of returning unpacked + train/eval data at each iteration + """ + + _target: Type = field(default_factory=lambda: DataManager) + """Target class to instantiate.""" + data: Optional[Path] = None + """Source of data, may not be used by all models.""" + camera_optimizer: Optional[CameraOptimizerConfig] = None + """Specifies the camera pose optimizer used during training. Helpful if poses are noisy.""" + + class DataManager(nn.Module): """Generic data manager's abstract class @@ -213,12 +229,10 @@ class DataManager(nn.Module): """Sets up the data manager for training. Here you will define any subclass specific object attributes from the attribute""" - raise NotImplementedError @abstractmethod def setup_eval(self): """Sets up the data manager for evaluation""" - raise NotImplementedError @abstractmethod def next_train(self, step: int) -> Tuple: @@ -241,6 +255,20 @@ class DataManager(nn.Module): """Returns the next eval image.""" raise NotImplementedError + @abstractmethod + def get_train_rays_per_batch(self) -> int: + """Returns the number of rays per batch for training.""" + raise NotImplementedError + + @abstractmethod + def get_eval_rays_per_batch(self) -> int: + """Returns the number of rays per batch for evaluation.""" + raise NotImplementedError + + def get_datapath(self) -> Optional[Path]: # pylint:disable=no-self-use + """Returns the path to the data. This is used to determine where to save camera paths.""" + return None + def get_training_callbacks( # pylint:disable=no-self-use self, training_callback_attributes: TrainingCallbackAttributes # pylint: disable=unused-argument ) -> List[TrainingCallback]: @@ -258,11 +286,8 @@ class DataManager(nn.Module): @dataclass -class VanillaDataManagerConfig(InstantiateConfig): - """Configuration for data manager instantiation; DataManager is in charge of keeping the train/eval dataparsers; - After instantiation, data manager holds both train/eval datasets and is in charge of returning unpacked - train/eval data at each iteration - """ +class VanillaDataManagerConfig(DataManagerConfig): + """A basic data manager""" _target: Type = field(default_factory=lambda: VanillaDataManager) """Target class to instantiate.""" @@ -333,7 +358,12 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method self.sampler = None self.test_mode = test_mode self.test_split = "test" if test_mode in ["test", "inference"] else "val" - self.dataparser = self.config.dataparser.setup() + self.dataparser = self.config.dataparser + if self.config.data is not None: + self.config.dataparser.data = Path(self.config.data) + else: + self.config.data = self.config.dataparser.data + self.dataparser = self.dataparser.setup() self.train_dataparser_outputs = self.dataparser.get_dataparser_outputs(split="train") self.train_dataset = self.create_train_dataset() @@ -455,6 +485,15 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method return image_idx, camera_ray_bundle, batch raise ValueError("No more eval images") + def get_train_rays_per_batch(self) -> int: + return self.config.train_num_rays_per_batch + + def get_eval_rays_per_batch(self) -> int: + return self.config.eval_num_rays_per_batch + + def get_datapath(self) -> Path: + return self.config.dataparser.data + def get_param_groups(self) -> Dict[str, List[Parameter]]: # pylint: disable=no-self-use """Get the param groups for the data manager. Returns: diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index f900dd863..4e38b4bef 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -126,20 +126,8 @@ class Trainer: # directory to save checkpoints self.checkpoint_dir: Path = config.get_checkpoint_dir() CONSOLE.log(f"Saving checkpoints to: {self.checkpoint_dir}") - # set up viewer if enabled - viewer_log_path: Path = self.base_dir / config.viewer.relative_log_filename - self.viewer_state, banner_messages = None, None - if self.config.is_viewer_enabled() and local_rank == 0: - self.viewer_state, banner_messages = viewer_utils.setup_viewer( - config.viewer, log_filename=viewer_log_path, datapath=config.pipeline.datamanager.dataparser.data - ) - self._check_viewer_warnings() - # set up writers/profilers if enabled - writer_log_path: Path = self.base_dir / config.logging.relative_log_dir - writer.setup_event_writer(config.is_wandb_enabled(), config.is_tensorboard_enabled(), log_dir=writer_log_path) - writer.setup_local_writer(config.logging, max_iter=config.max_num_iterations, banner_messages=banner_messages) - writer.put_config(name="config", config_dict=dataclasses.asdict(config), step=0) - profiler.setup_profiler(config.logging) + + self.viewer_state = None def setup(self, test_mode: Literal["test", "val", "inference"] = "val") -> None: """Setup the Trainer by calling other setup functions. @@ -165,6 +153,28 @@ class Trainer: ) ) + # set up viewer if enabled + viewer_log_path = self.base_dir / self.config.viewer.relative_log_filename + self.viewer_state, banner_messages = None, None + if self.config.is_viewer_enabled() and self.local_rank == 0: + datapath = self.pipeline.datamanager.get_datapath() + if datapath is None: + datapath = self.base_dir + self.viewer_state, banner_messages = viewer_utils.setup_viewer( + self.config.viewer, log_filename=viewer_log_path, datapath=datapath + ) + self._check_viewer_warnings() + # set up writers/profilers if enabled + writer_log_path = self.base_dir / self.config.logging.relative_log_dir + writer.setup_event_writer( + self.config.is_wandb_enabled(), self.config.is_tensorboard_enabled(), log_dir=writer_log_path + ) + writer.setup_local_writer( + self.config.logging, max_iter=self.config.max_num_iterations, banner_messages=banner_messages + ) + writer.put_config(name="config", config_dict=dataclasses.asdict(self.config), step=0) + profiler.setup_profiler(self.config.logging) + def setup_optimizers(self) -> Optimizers: """Helper to set up the optimizers @@ -172,9 +182,9 @@ class Trainer: The optimizers object given the trainer config. """ optimizer_config = self.config.optimizers.copy() - camera_optimizer_config = self.config.pipeline.datamanager.camera_optimizer param_groups = self.pipeline.get_param_groups() - if camera_optimizer_config.mode != "off": + camera_optimizer_config = self.config.pipeline.datamanager.camera_optimizer + if camera_optimizer_config is not None and camera_optimizer_config.mode != "off": assert camera_optimizer_config.param_group not in optimizer_config optimizer_config[camera_optimizer_config.param_group] = { "optimizer": camera_optimizer_config.optimizer, @@ -215,7 +225,7 @@ class Trainer: if step > 1: writer.put_time( name=EventName.TRAIN_RAYS_PER_SEC, - duration=self.config.pipeline.datamanager.train_num_rays_per_batch / train_t.duration, + duration=self.pipeline.datamanager.get_train_rays_per_batch() / train_t.duration, step=step, avg_over_steps=True, ) @@ -249,7 +259,7 @@ class Trainer: @check_main_thread def _always_render(self, step: int) -> None: - if self.config.is_viewer_enabled(): + if self.viewer_state is not None: while True: self.viewer_state.vis["renderingState/isTraining"].write(False) self._update_viewer_state(step) @@ -289,7 +299,7 @@ class Trainer: """ assert self.viewer_state is not None with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as _: - num_rays_per_batch: int = self.config.pipeline.datamanager.train_num_rays_per_batch + num_rays_per_batch: int = self.pipeline.datamanager.get_train_rays_per_batch() try: self.viewer_state.update_scene(self, step, self.pipeline.model, num_rays_per_batch) except RuntimeError: @@ -308,7 +318,7 @@ class Trainer: vis_t: timer object carrying time to execute visualization step step: current step """ - train_num_rays_per_batch: int = self.config.pipeline.datamanager.train_num_rays_per_batch + train_num_rays_per_batch: int = self.pipeline.datamanager.get_train_rays_per_batch() writer.put_time( name=EventName.TRAIN_RAYS_PER_SEC, duration=train_num_rays_per_batch / (train_t.duration - vis_t.duration), diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index c1cac0369..fd5d6fcc7 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -40,6 +40,7 @@ from typing_extensions import Literal from nerfstudio.configs import base_config as cfg from nerfstudio.data.datamanagers.base_datamanager import ( DataManager, + DataManagerConfig, VanillaDataManager, VanillaDataManagerConfig, ) @@ -193,7 +194,7 @@ class VanillaPipelineConfig(cfg.InstantiateConfig): _target: Type = field(default_factory=lambda: VanillaPipeline) """target class to instantiate""" - datamanager: VanillaDataManagerConfig = VanillaDataManagerConfig() + datamanager: DataManagerConfig = VanillaDataManagerConfig() """specifies the datamanager config""" model: ModelConfig = ModelConfig() """specifies the model config""" @@ -264,15 +265,16 @@ class VanillaPipeline(Pipeline): model_outputs = self.model(ray_bundle) metrics_dict = self.model.get_metrics_dict(model_outputs, batch) - camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group - if camera_opt_param_group in self.datamanager.get_param_groups(): - # Report the camera optimization metrics - metrics_dict["camera_opt_translation"] = ( - self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, :3].norm() - ) - metrics_dict["camera_opt_rotation"] = ( - self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, 3:].norm() - ) + if self.config.datamanager.camera_optimizer is not None: + camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group + if camera_opt_param_group in self.datamanager.get_param_groups(): + # Report the camera optimization metrics + metrics_dict["camera_opt_translation"] = ( + self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, :3].norm() + ) + metrics_dict["camera_opt_rotation"] = ( + self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, 3:].norm() + ) loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict) diff --git a/scripts/train.py b/scripts/train.py index f19b0ccc5..8f2e0b719 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -219,8 +219,8 @@ def main(config: TrainerConfig) -> None: config.set_timestamp() if config.data: - CONSOLE.log("Using --data alias for --data.pipeline.datamanager.dataparser.data") - config.pipeline.datamanager.dataparser.data = config.data + CONSOLE.log("Using --data alias for --data.pipeline.datamanager.data") + config.pipeline.datamanager.data = config.data if config.load_config: CONSOLE.log(f"Loading pre-set config from: {config.load_config}") diff --git a/scripts/viewer/run_viewer.py b/scripts/viewer/run_viewer.py index 9a96c6703..c48b23d9b 100644 --- a/scripts/viewer/run_viewer.py +++ b/scripts/viewer/run_viewer.py @@ -58,11 +58,11 @@ class RunViewer: self._start_viewer(config, pipeline) - def _start_viewer(self, config: TrainerConfig, pipeline): + def _start_viewer(self, config: TrainerConfig, pipeline: Pipeline): base_dir = config.get_base_dir() viewer_log_path = base_dir / config.viewer.relative_log_filename viewer_state, banner_messages = viewer_utils.setup_viewer( - config.viewer, log_filename=viewer_log_path, datapath=config.pipeline.datamanager.dataparser.data + config.viewer, log_filename=viewer_log_path, datapath=pipeline.datamanager.get_datapath() ) # We don't need logging, but writer.GLOBAL_BUFFER needs to be populated @@ -76,16 +76,16 @@ class RunViewer: ) while True: viewer_state.vis["renderingState/isTraining"].write(False) - self._update_viewer_state(viewer_state, config, pipeline) + self._update_viewer_state(viewer_state, pipeline) - def _update_viewer_state(self, viewer_state: viewer_utils.ViewerState, config: TrainerConfig, pipeline: Pipeline): + def _update_viewer_state(self, viewer_state: viewer_utils.ViewerState, pipeline: Pipeline): """Updates the viewer state by rendering out scene with current pipeline Returns the time taken to render scene. """ # NOTE: step must be > 0 otherwise the rendering would not happen step = 1 - num_rays_per_batch = config.pipeline.datamanager.train_num_rays_per_batch + num_rays_per_batch = pipeline.datamanager.get_train_rays_per_batch() with TimeWriter(writer, EventName.ITER_VIS_TIME) as _: try: viewer_state.update_scene(self, step, pipeline.model, num_rays_per_batch)