From a83d443dc3882173e4f10b9bb0efc6881f0e9e19 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Mon, 27 Feb 2023 12:36:01 -0800 Subject: [PATCH] Remove dataparser --- nerfstudio/configs/experiment_config.py | 6 ++- nerfstudio/configs/method_configs.py | 4 +- .../data/datamanagers/base_datamanager.py | 50 +++++++++++++++++-- .../datamanagers/dreamfusion_datamanager.py | 33 +++++++----- nerfstudio/engine/trainer.py | 50 +++++++++++-------- nerfstudio/models/base_model.py | 2 + nerfstudio/pipelines/base_pipeline.py | 22 ++++---- scripts/train.py | 8 ++- scripts/viewer/run_viewer.py | 6 +-- 9 files changed, 124 insertions(+), 57 deletions(-) diff --git a/nerfstudio/configs/experiment_config.py b/nerfstudio/configs/experiment_config.py index 8ebd77467..b73b9c55e 100644 --- a/nerfstudio/configs/experiment_config.py +++ b/nerfstudio/configs/experiment_config.py @@ -72,7 +72,9 @@ 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""" + prompt: Optional[str] = None + """Alias for --pipeline.model.prompt""" relative_model_dir: Path = Path("nerfstudio_models/") """Relative path to save all checkpoints.""" @@ -96,7 +98,7 @@ class ExperimentConfig(InstantiateConfig): def set_experiment_name(self) -> None: """Dynamically set the experiment name""" if self.experiment_name is None: - self.experiment_name = str(self.pipeline.datamanager.dataparser.data).replace("../", "").replace("/", "-") + self.experiment_name = str(self.pipeline.datamanager.data).replace("../", "").replace("/", "-") def get_base_dir(self) -> Path: """Retrieve the base directory to set relative paths""" diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index b4f7ec868..cf3fc20ce 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -338,6 +338,7 @@ method_configs["phototourism"] = TrainerConfig( method_configs["dreamfusion"] = TrainerConfig( method_name="dreamfusion", + experiment_name="", steps_per_eval_batch=50, steps_per_eval_image=50, steps_per_save=200, @@ -346,9 +347,6 @@ method_configs["dreamfusion"] = TrainerConfig( pipeline=VanillaPipelineConfig( generative=True, datamanager=DreamFusionDataManagerConfig( - dataparser=NerfstudioDataParserConfig(), - train_num_rays_per_batch=4096, - eval_num_rays_per_batch=4096, horizontal_rotation_warmup=4000, ), model=DreamFusionModelConfig( diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 0e3c6edce..db4ae8ebb 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 @@ -83,6 +84,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 @@ -201,12 +217,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: @@ -229,6 +243,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]: @@ -246,7 +274,7 @@ class DataManager(nn.Module): @dataclass -class VanillaDataManagerConfig(InstantiateConfig): +class VanillaDataManagerConfig(DataManagerConfig): """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 @@ -319,7 +347,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() @@ -438,6 +471,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/data/datamanagers/dreamfusion_datamanager.py b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py index 48a686df1..34780a1eb 100644 --- a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py +++ b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py @@ -18,6 +18,7 @@ Data manager for dreamfusion from __future__ import annotations +import sys from dataclasses import dataclass, field from typing import Dict, List, Tuple, Type, Union @@ -29,11 +30,7 @@ from typing_extensions import Literal from nerfstudio.cameras.cameras import Cameras from nerfstudio.cameras.rays import RayBundle from nerfstudio.configs.config_utils import to_immutable_dict -from nerfstudio.data.datamanagers.base_datamanager import ( - DataManager, - VanillaDataManager, - VanillaDataManagerConfig, -) +from nerfstudio.data.datamanagers.base_datamanager import DataManager, DataManagerConfig from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.scene_box import SceneBox @@ -146,11 +143,8 @@ def random_train_pose( @dataclass -class DreamFusionDataManagerConfig(VanillaDataManagerConfig): - """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 DreamFusionDataManagerConfig(DataManagerConfig): + """Configuration for data manager that does not load from a dataset. Instead, it generates random poses.""" _target: Type = field(default_factory=lambda: DreamFusionDataManager) train_resolution: int = 64 @@ -179,7 +173,7 @@ class DreamFusionDataManagerConfig(VanillaDataManagerConfig): """How many steps until the full horizontal rotation range is used""" -class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-method +class DreamFusionDataManager(DataManager): # pylint: disable=abstract-method """Basic stored data manager implementation. This is pretty much a port over from our old dataloading utilities, and is a little jank @@ -211,7 +205,10 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me 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() + + if self.config.data is not None: + CONSOLE.print("[red] --data should not be used with the DreamFusionDataManager[/red]") + sys.exit(1) cameras, _, _ = random_train_pose( self.config.num_eval_angles, @@ -266,7 +263,11 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me # camera_idx = torch.randint(0, self.eval_cameras.shape[0], [1], dtype=torch.long, device=self.device) # ray_bundle = self.eval_cameras.generate_rays(camera_idx).flatten() - return ray_bundle, {"vertical": vertical_rotation, "central": central_rotation, "initialization": True} + return ray_bundle, { + "vertical": vertical_rotation, + "central": central_rotation, + "initialization": True, + } def next_eval(self, step: int) -> Tuple[RayBundle, Dict]: """Returns the next batch of data from the eval dataloader.""" @@ -289,6 +290,12 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me return ray_bundle, {"vertical": vertical_rotation, "central": central_rotation} + def get_train_rays_per_batch(self) -> int: + return self.config.train_resolution**2 + + def get_eval_rays_per_batch(self) -> int: + return self.config.eval_resolution**2 + 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 5a29a6849..28c2f7b33 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -117,24 +117,12 @@ class Trainer: # optimizers self.grad_scaler = GradScaler(enabled=self.mixed_precision) + self.viewer_state = None + self.base_dir = config.get_base_dir() # directory to save checkpoints self.checkpoint_dir = config.get_checkpoint_dir() CONSOLE.log(f"Saving checkpoints to: {self.checkpoint_dir}") - # set up viewer if enabled - viewer_log_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 = 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) def setup(self, test_mode: Literal["test", "val", "inference"] = "val"): """Setup the Trainer by calling other setup functions. @@ -164,6 +152,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 @@ -171,9 +181,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): - 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 = self.config.pipeline.datamanager.train_num_rays_per_batch + num_rays_per_batch = 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 = self.config.pipeline.datamanager.train_num_rays_per_batch + train_num_rays_per_batch = 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/models/base_model.py b/nerfstudio/models/base_model.py index 3cc5fa2ec..58de64f95 100644 --- a/nerfstudio/models/base_model.py +++ b/nerfstudio/models/base_model.py @@ -50,6 +50,8 @@ class ModelConfig(InstantiateConfig): """parameters to instantiate density field with""" eval_num_rays_per_chunk: int = 4096 """specifies number of rays per chunk during eval""" + prompt: Optional[str] = None + """A prompt to be used in text to NeRF models""" class Model(nn.Module): diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index bf4d400f8..1ef5ffea5 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -41,6 +41,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, ) @@ -187,7 +188,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..bb3887133 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -219,8 +219,12 @@ 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.prompt: + CONSOLE.log("Using --prompt alias for --data.pipeline.model.prompt") + config.pipeline.model.prompt = config.prompt 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..336a424b1 100644 --- a/scripts/viewer/run_viewer.py +++ b/scripts/viewer/run_viewer.py @@ -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)