TrainerConfig inherit from InstantiateConfig (#1175)

Making TrainerConfig an InstantiateConfig so it can be swapped in and out with different methods. Also made the TrainerConfig a subclass of ExperimentConfig (which was changed from a printableconfig to an instantiateconfig), since we were already instantiating the Trainer class with the ExperimentConfig

Removed the "config" field from the TrainerCallbackAttributes struct to eliminate circular dependencies when types are checked.

Minor scripts and docs changes to accommodate the update to the config struct, meaning the flags that were associated with the TrainerConfig were moved up to the top level of the config.

nerfstudio/engine/callbacks.py

@@ -15,14 +15,13 @@
 """
 Callback code used for training iterations
 """
+from __future__ import annotations
 
 from dataclasses import InitVar, dataclass
 from enum import Enum, auto
 from inspect import signature
 from typing import Callable, Dict, List, Optional, Tuple
 
-from nerfstudio.configs.base_config import TrainerConfig
-
 
 @dataclass
 class TrainingCallbackAttributes:
@@ -38,8 +37,6 @@ class TrainingCallbackAttributes:
     """gradient scalers"""
     pipeline: Optional[InitVar]
     """reference to training pipeline"""
-    config: TrainerConfig
-    """the trainer config"""
 
 
 class TrainingCallbackLocation(Enum):

nerfstudio/engine/trainer.py

@@ -21,7 +21,9 @@ import dataclasses
 import functools
 import os
 import time
-from typing import Dict, List, Tuple
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Type
 
 import torch
 from rich.console import Console
@@ -49,6 +51,34 @@ from nerfstudio.viewer.server import viewer_utils
 CONSOLE = Console(width=120)
 
 
+@dataclass
+class TrainerConfig(ExperimentConfig):
+    """Configuration for training regimen"""
+
+    _target: Type = field(default_factory=lambda: Trainer)
+    """target class to instantiate"""
+    steps_per_save: int = 1000
+    """Number of steps between saves."""
+    steps_per_eval_batch: int = 500
+    """Number of steps between randomly sampled batches of rays."""
+    steps_per_eval_image: int = 500
+    """Number of steps between single eval images."""
+    steps_per_eval_all_images: int = 25000
+    """Number of steps between eval all images."""
+    max_num_iterations: int = 1000000
+    """Maximum number of iterations to run."""
+    mixed_precision: bool = False
+    """Whether or not to use mixed precision for training."""
+    save_only_latest_checkpoint: bool = True
+    """Whether to only save the latest checkpoint or all checkpoints."""
+    # optional parameters if we want to resume training
+    load_dir: Optional[Path] = None
+    """Optionally specify a pre-trained model directory to load from."""
+    load_step: Optional[int] = None
+    """Optionally specify model step to load from; if none, will find most recent model in load_dir."""
+    load_config: Optional[Path] = None
+
+
 class Trainer:
     """Trainer class
 
@@ -71,12 +101,12 @@ class Trainer:
     optimizers: Optimizers
     callbacks: List[TrainingCallback]
 
-    def __init__(self, config: ExperimentConfig, local_rank: int = 0, world_size: int = 1):
+    def __init__(self, config: TrainerConfig, local_rank: int = 0, world_size: int = 1):
         self.config = config
         self.local_rank = local_rank
         self.world_size = world_size
         self.device = "cpu" if world_size == 0 else f"cuda:{local_rank}"
-        self.mixed_precision = self.config.trainer.mixed_precision
+        self.mixed_precision = self.config.mixed_precision
         if self.device == "cpu":
             self.mixed_precision = False
             CONSOLE.print("Mixed precision is disabled for CPU training.")
@@ -97,9 +127,7 @@ class Trainer:
         # 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.trainer.max_num_iterations, banner_messages=banner_messages
-        )
+        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)
 
@@ -124,7 +152,6 @@ class Trainer:
                 optimizers=self.optimizers,  # type: ignore
                 grad_scaler=self.grad_scaler,  # type: ignore
                 pipeline=self.pipeline,  # type: ignore
-                config=self.config.trainer,  # type: ignore
             )
         )
 
@@ -132,7 +159,7 @@ class Trainer:
         """Helper to set up the optimizers
 
         Returns:
-            The optimizers object given the experiment config.
+            The optimizers object given the trainer config.
         """
         optimizer_config = self.config.optimizers.copy()
         camera_optimizer_config = self.config.pipeline.datamanager.camera_optimizer
@@ -155,7 +182,7 @@ class Trainer:
 
         self._init_viewer_state()
         with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME):
-            num_iterations = self.config.trainer.max_num_iterations
+            num_iterations = self.config.max_num_iterations
             step = 0
             for step in range(self._start_step, self._start_step + num_iterations):
                 with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as train_t:
@@ -194,7 +221,7 @@ class Trainer:
 
                 self.eval_iteration(step)
 
-                if step_check(step, self.config.trainer.steps_per_save):
+                if step_check(step, self.config.steps_per_save):
                     self.save_checkpoint(step)
 
                 writer.write_out_storage()
@@ -274,9 +301,9 @@ class Trainer:
 
     def _load_checkpoint(self) -> None:
         """Helper function to load pipeline and optimizer from prespecified checkpoint"""
-        load_dir = self.config.trainer.load_dir
+        load_dir = self.config.load_dir
         if load_dir is not None:
-            load_step = self.config.trainer.load_step
+            load_step = self.config.load_step
             if load_step is None:
                 print("Loading latest checkpoint from load_dir")
                 # NOTE: this is specific to the checkpoint name format
@@ -317,7 +344,7 @@ class Trainer:
             ckpt_path,
         )
         # possibly delete old checkpoints
-        if self.config.trainer.save_only_latest_checkpoint:
+        if self.config.save_only_latest_checkpoint:
             # delete everything else in the checkpoint folder
             for f in self.checkpoint_dir.glob("*"):
                 if f != ckpt_path:
@@ -352,7 +379,7 @@ class Trainer:
             step: Current training step.
         """
         # a batch of eval rays
-        if step_check(step, self.config.trainer.steps_per_eval_batch):
+        if step_check(step, self.config.steps_per_eval_batch):
             _, eval_loss_dict, eval_metrics_dict = self.pipeline.get_eval_loss_dict(step=step)
             eval_loss = functools.reduce(torch.add, eval_loss_dict.values())
             writer.put_scalar(name="Eval Loss", scalar=eval_loss, step=step)
@@ -360,7 +387,7 @@ class Trainer:
             writer.put_dict(name="Eval Metrics Dict", scalar_dict=eval_metrics_dict, step=step)
 
         # one eval image
-        if step_check(step, self.config.trainer.steps_per_eval_image):
+        if step_check(step, self.config.steps_per_eval_image):
             with TimeWriter(writer, EventName.TEST_RAYS_PER_SEC, write=False) as test_t:
                 metrics_dict, images_dict = self.pipeline.get_eval_image_metrics_and_images(step=step)
             writer.put_time(
@@ -375,6 +402,6 @@ class Trainer:
                 writer.put_image(name=group + "/" + image_name, image=image, step=step)
 
         # all eval images
-        if step_check(step, self.config.trainer.steps_per_eval_all_images):
+        if step_check(step, self.config.steps_per_eval_all_images):
             metrics_dict = self.pipeline.get_average_eval_image_metrics(step=step)
             writer.put_dict(name="Eval Images Metrics Dict (all images)", scalar_dict=metrics_dict, step=step)