mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
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.
This commit is contained in:
Jake Austin
GitHub
parent
b1dfd28cb6
commit
3ffee522e0
4
.vscode/launch.json
vendored
4
.vscode/launch.json
vendored
@ -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",
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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"""
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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}" \
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user