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