Wandb/experiment improvements (#1879)

exposing experiment settings

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

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

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:

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