From 3ffee522e07f22cb8c5df4bc3efe031fa44258e3 Mon Sep 17 00:00:00 2001 From: Jake Austin <55679937+jake-austin@users.noreply.github.com> Date: Tue, 27 Dec 2022 18:19:34 -1000 Subject: [PATCH] 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. --- .vscode/launch.json | 4 +- README.md | 2 +- docs/quickstart/first_nerf.md | 2 +- nerfstudio/cameras/camera_optimizers.py | 4 +- nerfstudio/configs/base_config.py | 29 ---------- nerfstudio/configs/experiment_config.py | 16 ++--- nerfstudio/configs/method_configs.py | 61 +++++++++++--------- nerfstudio/engine/callbacks.py | 5 +- nerfstudio/engine/trainer.py | 59 ++++++++++++++----- nerfstudio/utils/eval_utils.py | 11 ++-- scripts/benchmarking/launch_train_blender.sh | 6 +- scripts/train.py | 23 ++++---- tests/test_train.py | 6 +- 13 files changed, 114 insertions(+), 114 deletions(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index 8b92b869d..935d7ae8b 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -77,8 +77,8 @@ "--viewer.enable", "--viewer.zmq-url=tcp://127.0.0.1:6001", "--viewer.websocket-port=8014", - "--trainer.steps_per_save=1000000", - "--trainer.steps_per_test=1000000", + "--steps_per_save=1000000", + "--steps_per_test=1000000", "--pipeline.model.density-field-params.base-scale=4", "--pipeline.model.density-field-params.num-cascades=1", "pipeline.datamanager.train-dataparser:friends-data-parser-config", diff --git a/README.md b/README.md index edef31622..923753098 100644 --- a/README.md +++ b/README.md @@ -151,7 +151,7 @@ Navigating to the link at the end of the terminal will load the webviewer. If yo It is possible to load a pretrained model by running ```bash -ns-train nerfacto --data data/nerfstudio/poster --trainer.load-dir {outputs/.../nerfstudio_models} +ns-train nerfacto --data data/nerfstudio/poster --load-dir {outputs/.../nerfstudio_models} ``` This will automatically start training. If you do not want it to train, add `--viewer.start-train False` to your training command. diff --git a/docs/quickstart/first_nerf.md b/docs/quickstart/first_nerf.md index 1bd6ec367..30143814f 100644 --- a/docs/quickstart/first_nerf.md +++ b/docs/quickstart/first_nerf.md @@ -35,7 +35,7 @@ Navigating to the link at the end of the terminal will load the webviewer. If yo It is possible to load a pretrained model by running ```bash -ns-train nerfacto --data data/nerfstudio/poster --trainer.load-dir {outputs/.../nerfstudio_models} +ns-train nerfacto --data data/nerfstudio/poster --load-dir {outputs/.../nerfstudio_models} ``` This will automatically start training. If you do not want it to train, add `--viewer.start-train False` to your training command. diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 891cd272a..b8d846f02 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -29,14 +29,14 @@ from torchtyping import TensorType from typing_extensions import Literal, assert_never from nerfstudio.cameras.lie_groups import exp_map_SE3, exp_map_SO3xR3 -from nerfstudio.configs import base_config as cfg +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.engine.optimizers import AdamOptimizerConfig from nerfstudio.engine.schedulers import SchedulerConfig from nerfstudio.utils import poses as pose_utils @dataclass -class CameraOptimizerConfig(cfg.InstantiateConfig): +class CameraOptimizerConfig(InstantiateConfig): """Configuration of optimization for camera poses.""" _target: Type = field(default_factory=lambda: CameraOptimizer) diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index a74f07ed7..aa0177ce3 100644 --- a/nerfstudio/configs/base_config.py +++ b/nerfstudio/configs/base_config.py @@ -120,35 +120,6 @@ class LoggingConfig(PrintableConfig): profiler logs run times of functions and prints at end of training""" -# Trainer related configs -@dataclass -class TrainerConfig(PrintableConfig): - """Configuration for training regimen""" - - 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.""" - relative_model_dir: Path = Path("nerfstudio_models/") - """Relative path to save all checkpoints.""" - 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 - - # Viewer related configs @dataclass class ViewerConfig(PrintableConfig): diff --git a/nerfstudio/configs/experiment_config.py b/nerfstudio/configs/experiment_config.py index e78c8de59..16bbf129d 100644 --- a/nerfstudio/configs/experiment_config.py +++ b/nerfstudio/configs/experiment_config.py @@ -14,6 +14,8 @@ """Config used for running an experiment""" +from __future__ import annotations + from dataclasses import dataclass from datetime import datetime from pathlib import Path @@ -24,10 +26,9 @@ from rich.console import Console from typing_extensions import Literal from nerfstudio.configs.base_config import ( + InstantiateConfig, LoggingConfig, MachineConfig, - PrintableConfig, - TrainerConfig, ViewerConfig, ) from nerfstudio.configs.config_utils import to_immutable_dict @@ -39,8 +40,9 @@ CONSOLE = Console(width=120) @dataclass -class ExperimentConfig(PrintableConfig): - """Full config contents""" +class ExperimentConfig(InstantiateConfig): + """Full config contents for running an experiment. Any experiment types (like training) will be + subclassed from this, and must have their _target field defined accordingly.""" output_dir: Path = Path("outputs") """relative or absolute output directory to save all checkpoints and logging""" @@ -56,8 +58,6 @@ class ExperimentConfig(PrintableConfig): """Logging configuration""" viewer: ViewerConfig = ViewerConfig() """Viewer configuration""" - trainer: TrainerConfig = TrainerConfig() - """Trainer configuration""" pipeline: VanillaPipelineConfig = VanillaPipelineConfig() """Pipeline configuration""" optimizers: Dict[str, Any] = to_immutable_dict( @@ -73,6 +73,8 @@ class ExperimentConfig(PrintableConfig): """Which visualizer to use.""" data: Optional[Path] = None """Alias for --pipeline.datamanager.dataparser.data""" + relative_model_dir: Path = Path("nerfstudio_models/") + """Relative path to save all checkpoints.""" def is_viewer_enabled(self) -> bool: """Checks if a viewer is enabled.""" @@ -105,7 +107,7 @@ class ExperimentConfig(PrintableConfig): def get_checkpoint_dir(self) -> Path: """Retrieve the checkpoint directory""" - return Path(self.get_base_dir() / self.trainer.relative_model_dir) + return Path(self.get_base_dir() / self.relative_model_dir) def print_to_terminal(self) -> None: """Helper to pretty print config to terminal""" diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 3d4b694d9..e7995cd71 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -24,8 +24,7 @@ import tyro from nerfacc import ContractionType from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig -from nerfstudio.configs.base_config import TrainerConfig, ViewerConfig -from nerfstudio.configs.experiment_config import ExperimentConfig +from nerfstudio.configs.base_config import ViewerConfig from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConfig from nerfstudio.data.datamanagers.semantic_datamanager import SemanticDataManagerConfig from nerfstudio.data.datamanagers.variable_res_datamanager import ( @@ -43,6 +42,7 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import ( ) from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig from nerfstudio.engine.schedulers import SchedulerConfig +from nerfstudio.engine.trainer import TrainerConfig from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind from nerfstudio.models.instant_ngp import InstantNGPModelConfig from nerfstudio.models.mipnerf import MipNerfModel @@ -53,7 +53,7 @@ from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig from nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig -method_configs: Dict[str, ExperimentConfig] = {} +method_configs: Dict[str, TrainerConfig] = {} descriptions = { "nerfacto": "Recommended real-time model tuned for real captures. This model will be continually updated.", "instant-ngp": "Implementation of Instant-NGP. Recommended real-time model for unbounded scenes.", @@ -66,11 +66,12 @@ descriptions = { "phototourism": "Uses the Phototourism data.", } -method_configs["nerfacto"] = ExperimentConfig( +method_configs["nerfacto"] = TrainerConfig( method_name="nerfacto", - trainer=TrainerConfig( - steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True - ), + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=30000, + mixed_precision=True, pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( dataparser=NerfstudioDataParserConfig(), @@ -96,11 +97,12 @@ method_configs["nerfacto"] = ExperimentConfig( vis="viewer", ) -method_configs["instant-ngp"] = ExperimentConfig( +method_configs["instant-ngp"] = TrainerConfig( method_name="instant-ngp", - trainer=TrainerConfig( - steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True - ), + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=30000, + mixed_precision=True, pipeline=DynamicBatchPipelineConfig( datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=8192), model=InstantNGPModelConfig(eval_num_rays_per_chunk=8192), @@ -116,11 +118,12 @@ method_configs["instant-ngp"] = ExperimentConfig( ) -method_configs["instant-ngp-bounded"] = ExperimentConfig( +method_configs["instant-ngp-bounded"] = TrainerConfig( method_name="instant-ngp-bounded", - trainer=TrainerConfig( - steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True - ), + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=30000, + mixed_precision=True, pipeline=DynamicBatchPipelineConfig( datamanager=VanillaDataManagerConfig(dataparser=InstantNGPDataParserConfig(), train_num_rays_per_batch=8192), model=InstantNGPModelConfig( @@ -143,7 +146,7 @@ method_configs["instant-ngp-bounded"] = ExperimentConfig( ) -method_configs["mipnerf"] = ExperimentConfig( +method_configs["mipnerf"] = TrainerConfig( method_name="mipnerf", pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=1024), @@ -163,11 +166,12 @@ method_configs["mipnerf"] = ExperimentConfig( }, ) -method_configs["semantic-nerfw"] = ExperimentConfig( +method_configs["semantic-nerfw"] = TrainerConfig( method_name="semantic-nerfw", - trainer=TrainerConfig( - steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True - ), + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=30000, + mixed_precision=True, pipeline=VanillaPipelineConfig( datamanager=SemanticDataManagerConfig( dataparser=FriendsDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=8192 @@ -188,7 +192,7 @@ method_configs["semantic-nerfw"] = ExperimentConfig( vis="viewer", ) -method_configs["vanilla-nerf"] = ExperimentConfig( +method_configs["vanilla-nerf"] = TrainerConfig( method_name="vanilla-nerf", pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( @@ -208,9 +212,9 @@ method_configs["vanilla-nerf"] = ExperimentConfig( }, ) -method_configs["tensorf"] = ExperimentConfig( +method_configs["tensorf"] = TrainerConfig( method_name="tensorf", - trainer=TrainerConfig(mixed_precision=False), + mixed_precision=False, pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( dataparser=BlenderDataParserConfig(), @@ -229,7 +233,7 @@ method_configs["tensorf"] = ExperimentConfig( }, ) -method_configs["dnerf"] = ExperimentConfig( +method_configs["dnerf"] = TrainerConfig( method_name="dnerf", pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig(dataparser=DNeRFDataParserConfig()), @@ -251,11 +255,12 @@ method_configs["dnerf"] = ExperimentConfig( }, ) -method_configs["phototourism"] = ExperimentConfig( +method_configs["phototourism"] = TrainerConfig( method_name="phototourism", - trainer=TrainerConfig( - steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True - ), + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=30000, + mixed_precision=True, pipeline=VanillaPipelineConfig( datamanager=VariableResDataManagerConfig( # NOTE: one of the only differences with nerfacto dataparser=PhototourismDataParserConfig(), # NOTE: one of the only differences with nerfacto diff --git a/nerfstudio/engine/callbacks.py b/nerfstudio/engine/callbacks.py index 8f592669f..25f73d8d2 100644 --- a/nerfstudio/engine/callbacks.py +++ b/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): diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 513fb194e..68c1182c8 100644 --- a/nerfstudio/engine/trainer.py +++ b/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) diff --git a/nerfstudio/utils/eval_utils.py b/nerfstudio/utils/eval_utils.py index 8d780cc82..40f9e212e 100644 --- a/nerfstudio/utils/eval_utils.py +++ b/nerfstudio/utils/eval_utils.py @@ -27,8 +27,7 @@ import yaml from rich.console import Console from typing_extensions import Literal -from nerfstudio.configs.base_config import TrainerConfig -from nerfstudio.configs.experiment_config import ExperimentConfig +from nerfstudio.engine.trainer import TrainerConfig from nerfstudio.pipelines.base_pipeline import Pipeline CONSOLE = Console(width=120) @@ -69,7 +68,7 @@ def eval_setup( config_path: Path, eval_num_rays_per_chunk: Optional[int] = None, test_mode: Literal["test", "val", "inference"] = "test", -) -> Tuple[ExperimentConfig, Pipeline, Path]: +) -> Tuple[TrainerConfig, Pipeline, Path]: """Shared setup for loading a saved pipeline for evaluation. Args: @@ -86,14 +85,14 @@ def eval_setup( """ # load save config config = yaml.load(config_path.read_text(), Loader=yaml.Loader) - assert isinstance(config, ExperimentConfig) + assert isinstance(config, TrainerConfig) if eval_num_rays_per_chunk: config.pipeline.model.eval_num_rays_per_chunk = eval_num_rays_per_chunk # load checkpoints from wherever they were saved # TODO: expose the ability to choose an arbitrary checkpoint - config.trainer.load_dir = config.get_checkpoint_dir() + config.load_dir = config.get_checkpoint_dir() config.pipeline.datamanager.eval_image_indices = None # setup pipeline (which includes the DataManager) @@ -103,6 +102,6 @@ def eval_setup( pipeline.eval() # load checkpointed information - checkpoint_path = eval_load_checkpoint(config.trainer, pipeline) + checkpoint_path = eval_load_checkpoint(config, pipeline) return config, pipeline, checkpoint_path diff --git a/scripts/benchmarking/launch_train_blender.sh b/scripts/benchmarking/launch_train_blender.sh index 332ec95ef..2ddb0b008 100755 --- a/scripts/benchmarking/launch_train_blender.sh +++ b/scripts/benchmarking/launch_train_blender.sh @@ -79,9 +79,9 @@ for dataset in "${DATASETS[@]}"; do ns-train "${method_name}" "${method_opts[@]}" \ --data="data/blender/${dataset}${trans_file}" \ --experiment-name="blender_${dataset}_${tag}" \ - --trainer.relative-model-dir=nerfstudio_models/ \ - --trainer.steps-per-save=1000 \ - --trainer.max-num-iterations=16500 \ + --relative-model-dir=nerfstudio_models/ \ + --steps-per-save=1000 \ + --max-num-iterations=16500 \ --logging.local-writer.enable=False \ --logging.enable-profiler=False \ --vis "${vis}" \ diff --git a/scripts/train.py b/scripts/train.py index 910dbff1c..240fdeca0 100755 --- a/scripts/train.py +++ b/scripts/train.py @@ -46,9 +46,8 @@ import yaml from rich.console import Console from nerfstudio.configs.config_utils import convert_markup_to_ansi -from nerfstudio.configs.experiment_config import ExperimentConfig from nerfstudio.configs.method_configs import AnnotatedBaseConfigUnion -from nerfstudio.engine.trainer import Trainer +from nerfstudio.engine.trainer import TrainerConfig from nerfstudio.utils import comms, profiler CONSOLE = Console(width=120) @@ -74,7 +73,7 @@ def _set_random_seed(seed) -> None: torch.manual_seed(seed) -def train_loop(local_rank: int, world_size: int, config: ExperimentConfig, global_rank: int = 0): +def train_loop(local_rank: int, world_size: int, config: TrainerConfig, global_rank: int = 0): """Main training function that sets up and runs the trainer per process Args: @@ -83,7 +82,7 @@ def train_loop(local_rank: int, world_size: int, config: ExperimentConfig, globa config: config file specifying training regimen """ _set_random_seed(config.machine.seed + global_rank) - trainer = Trainer(config, local_rank, world_size) + trainer = config.setup(local_rank=local_rank, world_size=world_size) trainer.setup() trainer.train() @@ -95,7 +94,7 @@ def _distributed_worker( num_gpus_per_machine: int, machine_rank: int, dist_url: str, - config: ExperimentConfig, + config: TrainerConfig, timeout: timedelta = DEFAULT_TIMEOUT, ) -> Any: """Spawned distributed worker that handles the initialization of process group and handles the @@ -110,7 +109,7 @@ def _distributed_worker( dist_url: URL to connect to for distributed jobs, including protocol E.g., "tcp://127.0.0.1:8686". It can be set to "auto" to automatically select a free port on localhost. - config: ExperimentConfig specifying training regimen. + config: TrainerConfig specifying training regimen. timeout: Timeout of the distributed workers. Raises: @@ -150,7 +149,7 @@ def launch( num_machines: int = 1, machine_rank: int = 0, dist_url: str = "auto", - config: Optional[ExperimentConfig] = None, + config: Optional[TrainerConfig] = None, timeout: timedelta = DEFAULT_TIMEOUT, ) -> None: """Function that spawns muliple processes to call on main_func @@ -161,7 +160,7 @@ def launch( num_machines (int, optional): total number of machines machine_rank (int, optional): rank of this machine. dist_url (str, optional): url to connect to for distributed jobs. - config (ExperimentConfig, optional): config file specifying training regimen. + config (TrainerConfig, optional): config file specifying training regimen. timeout (timedelta, optional): timeout of the distributed workers. """ assert config is not None @@ -215,7 +214,7 @@ def launch( profiler.flush_profiler(config.logging) -def main(config: ExperimentConfig) -> None: +def main(config: TrainerConfig) -> None: """Main function.""" config.set_timestamp() @@ -223,9 +222,9 @@ def main(config: ExperimentConfig) -> None: CONSOLE.log("Using --data alias for --data.pipeline.datamanager.dataparser.data") config.pipeline.datamanager.dataparser.data = config.data - if config.trainer.load_config: - CONSOLE.log(f"Loading pre-set config from: {config.trainer.load_config}") - config = yaml.load(config.trainer.load_config.read_text(), Loader=yaml.Loader) + if config.load_config: + CONSOLE.log(f"Loading pre-set config from: {config.load_config}") + config = yaml.load(config.load_config.read_text(), Loader=yaml.Loader) # print and save config config.print_to_terminal() diff --git a/tests/test_train.py b/tests/test_train.py index 8b564369b..63a8099e9 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -9,19 +9,19 @@ from pathlib import Path import pytest -from nerfstudio.configs.experiment_config import ExperimentConfig from nerfstudio.configs.method_configs import method_configs from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.minimal_dataparser import MinimalDataParserConfig +from nerfstudio.engine.trainer import TrainerConfig from scripts.train import train_loop BLACKLIST = ["base", "semantic-nerfw", "instant-ngp", "instant-ngp-bounded", "nerfacto", "phototourism"] -def set_reduced_config(config: ExperimentConfig): +def set_reduced_config(config: TrainerConfig): """Reducing the config settings to speedup test""" config.machine.num_gpus = 0 - config.trainer.max_num_iterations = 2 + config.max_num_iterations = 2 # reduce dataset factors; set dataset to test config.pipeline.datamanager.dataparser = BlenderDataParserConfig(data=Path("tests/data/lego_test")) config.pipeline.datamanager.train_num_images_to_sample_from = 1