diff --git a/nerfstudio/configs/experiment_config.py b/nerfstudio/configs/experiment_config.py index 72362befe..c0da16fef 100644 --- a/nerfstudio/configs/experiment_config.py +++ b/nerfstudio/configs/experiment_config.py @@ -50,6 +50,8 @@ class ExperimentConfig(InstantiateConfig): """Method name. Required to set in python or via cli""" experiment_name: Optional[str] = None """Experiment name. If None, will automatically be set to dataset name""" + project_name: Optional[str] = "nerfstudio-project" + """Project name.""" timestamp: str = "{timestamp}" """Experiment timestamp.""" machine: MachineConfig = MachineConfig() diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index 23d712d7e..7f75b782a 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/nerfstudio/data/utils/dataloaders.py @@ -211,7 +211,6 @@ class FixedIndicesEvalDataloader(EvalDataloader): self.count = 0 def __iter__(self): - self.count = 0 return self def __next__(self): @@ -225,30 +224,16 @@ class FixedIndicesEvalDataloader(EvalDataloader): class RandIndicesEvalDataloader(EvalDataloader): """Dataloader that returns random images. - Args: input_dataset: InputDataset to load data from device: Device to load data to """ - def __init__( - self, - input_dataset: InputDataset, - device: Union[torch.device, str] = "cpu", - **kwargs, - ): - super().__init__(input_dataset, device, **kwargs) - self.count = 0 - def __iter__(self): - self.count = 0 return self def __next__(self): - if self.count < 1: - image_indices = range(self.cameras.size) - image_idx = random.choice(image_indices) - ray_bundle, batch = self.get_data_from_image_idx(image_idx) - self.count += 1 - return ray_bundle, batch - raise StopIteration + # choose a random image index + image_idx = random.randint(0, len(self.cameras) - 1) + ray_bundle, batch = self.get_data_from_image_idx(image_idx) + return ray_bundle, batch diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index afaa03655..5680bddbe 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -89,6 +89,8 @@ class TrainerConfig(ExperimentConfig): """Optionally specify model step to load from; if none, will find most recent model in load_dir.""" load_config: Optional[Path] = None """Path to config YAML file.""" + load_checkpoint: Optional[Path] = None + """Path to checkpoint file.""" log_gradients: bool = False """Optionally log gradients during training""" @@ -185,7 +187,11 @@ class Trainer: # 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 + self.config.is_wandb_enabled(), + self.config.is_tensorboard_enabled(), + log_dir=writer_log_path, + experiment_name=self.config.experiment_name, + project_name=self.config.project_name, ) writer.setup_local_writer( self.config.logging, max_iter=self.config.max_num_iterations, banner_messages=banner_messages @@ -370,10 +376,11 @@ class Trainer: def _load_checkpoint(self) -> None: """Helper function to load pipeline and optimizer from prespecified checkpoint""" load_dir: Path = self.config.load_dir + load_checkpoint: Path = self.config.load_checkpoint if load_dir is not None: load_step = self.config.load_step if load_step is None: - print("Loading latest checkpoint from load_dir") + print("Loading latest Nerfstudio checkpoint from load_dir...") # NOTE: this is specific to the checkpoint name format load_step = sorted(int(x[x.find("-") + 1 : x.find(".")]) for x in os.listdir(load_dir))[-1] load_path: Path = load_dir / f"step-{load_step:09d}.ckpt" @@ -384,9 +391,18 @@ class Trainer: self.pipeline.load_pipeline(loaded_state["pipeline"], loaded_state["step"]) self.optimizers.load_optimizers(loaded_state["optimizers"]) self.grad_scaler.load_state_dict(loaded_state["scalers"]) - CONSOLE.print(f"done loading checkpoint from {load_path}") + CONSOLE.print(f"Done loading Nerfstudio checkpoint from {load_path}") + elif load_checkpoint is not None: + assert load_checkpoint.exists(), f"Checkpoint {load_checkpoint} does not exist" + loaded_state = torch.load(load_checkpoint, map_location="cpu") + self._start_step = loaded_state["step"] + 1 + # load the checkpoints for pipeline, optimizers, and gradient scalar + self.pipeline.load_pipeline(loaded_state["pipeline"], loaded_state["step"]) + self.optimizers.load_optimizers(loaded_state["optimizers"]) + self.grad_scaler.load_state_dict(loaded_state["scalers"]) + CONSOLE.print(f"Done loading Nerfstudio checkpoint from {load_checkpoint}") else: - CONSOLE.print("No checkpoints to load, training from scratch") + CONSOLE.print("No Nerfstudio checkpoint to load, so training from scratch.") @check_main_thread def save_checkpoint(self, step: int) -> None: diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index 4d168d508..2036e4659 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -18,7 +18,6 @@ Generic Writer class from __future__ import annotations import enum -import os from abc import abstractmethod from pathlib import Path from time import time @@ -188,9 +187,14 @@ def setup_local_writer(config: cfg.LoggingConfig, max_iter: int, banner_messages @check_main_thread -def setup_event_writer(is_wandb_enabled: bool, is_tensorboard_enabled: bool, log_dir: Path) -> None: +def setup_event_writer( + is_wandb_enabled: bool, + is_tensorboard_enabled: bool, + log_dir: Path, + experiment_name: str, + project_name: str = "nerfstudio-project", +) -> None: """Initialization of all event writers specified in config - Args: config: configuration to instantiate loggers max_iter: maximum number of train iterations @@ -198,7 +202,7 @@ def setup_event_writer(is_wandb_enabled: bool, is_tensorboard_enabled: bool, log """ using_event_writer = False if is_wandb_enabled: - curr_writer = WandbWriter(log_dir=log_dir) + curr_writer = WandbWriter(log_dir=log_dir, experiment_name=experiment_name, project_name=project_name) EVENT_WRITERS.append(curr_writer) using_event_writer = True if is_tensorboard_enabled: @@ -282,13 +286,8 @@ class TimeWriter: class WandbWriter(Writer): """WandDB Writer Class""" - def __init__(self, log_dir: Path): - wandb.init( - project=os.environ.get("WANDB_PROJECT", "nerfstudio-project"), - dir=os.environ.get("WANDB_DIR", str(log_dir)), - name=os.environ.get("WANDB_NAME", log_dir.name), - reinit=True, - ) + def __init__(self, log_dir: Path, experiment_name: str, project_name: str = "nerfstudio-project"): + wandb.init(project=project_name, dir=str(log_dir), reinit=True, name=experiment_name) def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: image = torch.permute(image, (2, 0, 1))