# Copyright 2022 The Plenoptix Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Base Configs""" from __future__ import annotations from dataclasses import dataclass from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union import dcargs import torch from nerfactory.configs.utils import to_immutable_dict # data instances from nerfactory.datamanagers.base import VanillaDataManager from nerfactory.datamanagers.dataparsers.base import DataParser from nerfactory.datamanagers.dataparsers.blender_parser import Blender from nerfactory.datamanagers.dataparsers.friends_parser import Friends from nerfactory.datamanagers.dataparsers.instant_ngp_parser import InstantNGP from nerfactory.datamanagers.dataparsers.mipnerf_parser import Mipnerf360 from nerfactory.datamanagers.dataparsers.nerfactory_parser import Nerfactory from nerfactory.datamanagers.dataparsers.record3d_parser import Record3D # model instances from nerfactory.models.base import Model from nerfactory.models.compound import CompoundModel from nerfactory.models.instant_ngp import NGPModel from nerfactory.models.nerfw import NerfWModel from nerfactory.models.tensorf import TensoRFModel from nerfactory.optimizers.schedulers import ExponentialDecaySchedule from nerfactory.pipelines.base import Pipeline from nerfactory.utils import writer # Pretty printing class class PrintableConfig: # pylint: disable=too-few-public-methods """Printable Config defining str function""" def __str__(self): lines = [self.__class__.__name__ + ":"] for key, val in vars(self).items(): if isinstance(val, Tuple): flattened_val = "[" for item in val: flattened_val += str(item) + "\n" flattened_val = flattened_val.rstrip("\n") val = flattened_val + "]" lines += f"{key}: {str(val)}".split("\n") return "\n ".join(lines) # Base instantiate configs @dataclass class InstantiateConfig(PrintableConfig): # pylint: disable=too-few-public-methods """Config class for instantiating an the class specified in the _target attribute.""" _target: Type def setup(self, **kwargs) -> Any: """Returns the instantiated object using the config.""" return self._target(self, **kwargs) # Machine related configs @dataclass class MachineConfig(PrintableConfig): """Configuration of machine setup""" seed: int = 42 """random seed initilization""" num_gpus: int = 1 """total number of gpus available for train/eval""" num_machines: int = 1 """total number of distributed machines available (for DDP)""" machine_rank: int = 0 """current machine's rank (for DDP)""" dist_url: str = "auto" """distributed connection point (for DDP)""" # Logging related configs @dataclass class TensorboardWriterConfig(InstantiateConfig): """Tensorboard Writer config""" _target: Type = writer.TensorboardWriter """target class to instantiate""" enable: bool = False """if True enables tensorboard logging, else disables""" relative_log_dir: Path = Path("./") """relative path to save all tensorboard events""" log_dir: Optional[Path] = None # full log dir path to be dynamically set """auto populated absolute path to saved tensorboard events [Do not set!]""" @dataclass class WandbWriterConfig(InstantiateConfig): """WandDB Writer config""" _target: Type = writer.WandbWriter """target class to instantiate""" enable: bool = False """if True enables wandb logging, else disables""" relative_log_dir: Path = Path("./") """relative path to save all wandb events""" log_dir: Optional[Path] = None # full log dir path to be dynamically set """auto populated absolute path to saved wandb events [Do not set!]""" @dataclass class LocalWriterConfig(InstantiateConfig): """Local Writer config""" _target: Type = writer.LocalWriter """target class to instantiate""" enable: bool = False """if True enables local logging, else disables""" stats_to_track: Tuple[writer.EventName, ...] = ( writer.EventName.ITER_TRAIN_TIME, writer.EventName.TRAIN_RAYS_PER_SEC, writer.EventName.CURR_TEST_PSNR, writer.EventName.VIS_RAYS_PER_SEC, writer.EventName.TEST_RAYS_PER_SEC, ) """specifies which stats will be logged/printed to terminal""" max_log_size: int = 10 """maximum number of rows to print before wrapping. if 0, will print everything.""" relative_log_dir: Path = Path("./") """relative local path to save all events""" log_dir: Optional[Path] = None # full log dir path to be dynamically set """auto populated absolute local path to saved events [Do not set!]""" def setup(self, banner_messages: Optional[List[str]] = None, **kwargs) -> Any: """Instantiate local writer Args: banner_messages: List of strings that always print at the bottom of screen. Defaults to None. """ return self._target(self, banner_messages=banner_messages, **kwargs) @dataclass class LoggingConfig(PrintableConfig): """Configuration of loggers and profilers""" steps_per_log: int = 10 """number of steps between logging stats""" max_buffer_size: int = 20 """maximum history size to keep for computing running averages of stats. e.g. if 20, averages will be computed over past 20 occurances.""" writer: Tuple[Any, ...] = ( TensorboardWriterConfig(enable=True), WandbWriterConfig(enable=False), LocalWriterConfig(enable=True), ) """list of all supported writers. Can turn on/off writers by specifying enable.""" enable_profiler: bool = True """whether to enable profiling code; prints speed of functions at the end of a program. 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_test: int = 500 """number of steps between eval""" 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("nerfactory_models/") """relative path to save all checkpoints""" model_dir: Optional[Path] = None # full model dir path to be dynamically set """auto populated absolute path to saved checkpoints [Do not set!]""" # 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 """optionally specify a pre-defined config to load from""" # Dataset related configs @dataclass class DataParserConfig(InstantiateConfig): """Basic dataset config""" _target: Type = DataParser """_target: target class to instantiate""" @dataclass class NerfactoryDataParserConfig(DataParserConfig): """Nerfactory dataset config""" _target: Type = Nerfactory """target class to instantiate""" data_directory: Path = Path("data/ours/posterv2") """directory specifying location of data""" scale_factor: float = 1.0 """How much to scale the camera origins by.""" downscale_factor: int = 1 """How much to downscale images. Defaults to 1.""" scene_scale: float = 4.0 """How much to scale the scene. Defaults to 0.33""" orientation_method: Literal["pca", "up"] = "up" """The method to use for orientation. Either "pca" or "up".""" @dataclass class BlenderDataParserConfig(DataParserConfig): """Blender dataset config""" _target: Type = Blender """target class to instantiate""" data_directory: Path = Path("data/blender/lego") """directory specifying location of data""" scale_factor: float = 1.0 """How much to scale the camera origins by.""" alpha_color: str = "white" """alpha color of background""" @dataclass class FriendsDataParserConfig(DataParserConfig): """Friends dataset config""" _target: Type = Friends """target class to instantiate""" data_directory: Path = Path("data/friends/TBBT-big_living_room") """directory specifying location of data""" include_semantics: bool = True """whether or not to include loading of semantics data""" downscale_factor: int = 8 scene_scale: float = 4.0 """ Sets the bounding cube to have edge length of this size. The longest dimension of the Friends axis-aligned bbox will be scaled to this value. """ @dataclass class MipNerf360DataParserConfig(DataParserConfig): """Mipnerf 360 dataset config""" _target: Type = Mipnerf360 """target class to instantiate""" data_directory: Path = Path("data/mipnerf_360/garden") """directory specifying location of data""" downscale_factor: int = 1 """How much to downscale images. Defaults to 1.""" val_skip: int = 8 """1/val_skip images to use for validation. Defaults to 8.""" auto_scale: bool = True """Scale based on pose bounds. Defaults to True.""" aabb_scale: float = 4 """Scene scale, Defaults to 1.0.""" @dataclass class InstantNGPDataParserConfig(DataParserConfig): """Instant-NGP dataset config""" _target: Type = InstantNGP """target class to instantiate""" data_directory: Path = Path("data/ours/posterv2") """directory specifying location of data""" scale_factor: float = 1.0 """How much to scale the camera origins by.""" scene_scale: float = 0.33 """How much to scale the scene. Defaults to 0.33""" @dataclass class Record3DDataParserConfig(DataParserConfig): """Mipnerf 360 dataset config""" _target: Type = Record3D """target class to instantiate""" data_directory: Path = Path("data/record3d/garden") """Location of data""" val_skip: int = 8 """1/val_skip images to use for validation. Defaults to 8.""" aabb_scale: float = 4.0 """Scene scale, Defaults to 4.0.""" max_dataset_size: int = 150 """Max number of images to train on. If the dataset has more, images will be sampled approximately evenly. Defaults to 150.""" @dataclass class VanillaDataManagerConfig(InstantiateConfig): """Configuration for data manager instantiation; DataManager is in charge of keeping the train/eval dataparsers; After instantiation, data manager holds both train/eval datasets and is in charge of returning unpacked train/eval data at each iteration """ # Note: eval_dataparser is annotated with Fixed[] to prevent dcargs from trying to # convert Optional[InstantiateConfig] into subcommands for choosing between None and # InstantiateConfig. _target: Type = VanillaDataManager """target class to instantiate""" train_dataparser: Union[ NerfactoryDataParserConfig, BlenderDataParserConfig, FriendsDataParserConfig, MipNerf360DataParserConfig, InstantNGPDataParserConfig, Record3DDataParserConfig, ] = BlenderDataParserConfig() """specifies the dataparser used to unpack the data""" train_num_rays_per_batch: int = 1024 """number of rays per batch to use per training iteration""" train_num_images_to_sample_from: int = -1 """number of images to sample during training iteration""" eval_dataparser: dcargs.conf.Fixed[Optional[InstantiateConfig]] = None """optionally specify different dataparser to use during eval; if None, uses train_dataparser""" eval_image_indices: Optional[Tuple[int, ...]] = (0,) """specifies the image indices to use during eval; if None, uses all""" eval_num_rays_per_chunk: int = 4096 """specifies number of rays per chunk during eval""" @dataclass class FriendsDataManagerConfig(VanillaDataManagerConfig): """Friends data manager config""" _target: Type = VanillaDataManager """target class to instantiate""" train_dataparser: DataParserConfig = FriendsDataParserConfig() """specifies the dataparser used to unpack the data""" # Model related configs @dataclass class ModelConfig(InstantiateConfig): """Configuration for model instantiation""" _target: Type = Model """target class to instantiate""" enable_collider: bool = True """Whether to create a scene collider to filter rays.""" collider_params: Dict[str, float] = to_immutable_dict({"near_plane": 2.0, "far_plane": 6.0}) """parameters to instantiate scene collider with""" loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss_coarse": 1.0, "rgb_loss_fine": 1.0}) """Loss specific weights.""" num_coarse_samples: int = 64 """Number of samples in coarse field evaluation. Defaults to 64""" num_importance_samples: int = 128 """Number of samples in fine field evaluation. Defaults to 128""" field_implementation: Literal["torch", "tcnn"] = "torch" """one of "torch" or "tcnn", or other fields in 'field_implementation_to_class""" enable_density_field: bool = False """Whether to create a density field to filter samples.""" density_field_params: Dict[str, Any] = to_immutable_dict( { "center": 0.0, # simply set it as the center of the scene bbox "base_scale": 3.0, # simply set it as the scale of the scene bbox "num_cascades": 1, # if using more than 1 cascade, the `base_scale` can be smaller than scene scale. "resolution": 128, "update_every_num_iters": 16, } ) """parameters to instantiate density field with""" @dataclass class InstantNGPModelConfig(ModelConfig): """Instant NGP Model Config""" _target: Type = NGPModel """target class to instantiate""" enable_collider: bool = False """Whether to create a scene collider to filter rays.""" loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss": 1.0}) """Loss specific weights.""" field_implementation: Literal["torch", "tcnn"] = "tcnn" # torch, tcnn, ... """one of "torch" or "tcnn", or other fields in 'field_implementation_to_class'""" enable_density_field: bool = True """Whether to create a density field to filter samples.""" num_samples: int = 1024 # instead of course/fine samples """Number of samples in field evaluation. Defaults to 1024,""" @dataclass class CompoundModelConfig(ModelConfig): """Compound Model Config""" _target: Type = CompoundModel enable_density_field: bool = True enable_collider: bool = False field_implementation: Literal["torch", "tcnn"] = "tcnn" # torch, tcnn, ... loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss": 1.0}) num_samples: int = 1024 # instead of course/fine samples @dataclass class NerfWModelConfig(ModelConfig): """NerfW model config""" _target: Type = NerfWModel """target class to instantiate""" loss_coefficients: Dict[str, float] = to_immutable_dict( {"rgb_loss_coarse": 1.0, "rgb_loss_fine": 1.0, "uncertainty_loss": 1.0, "density_loss": 0.01} ) """Loss specific weights.""" num_coarse_samples: int = 64 """Number of samples in coarse field evaluation. Defaults to 64,""" num_importance_samples: int = 64 """Number of samples in fine field evaluation. Defaults to 128,""" uncertainty_min: float = 0.03 """This is added to the end of the uncertainty rendering operation. It's called 'beta_min' in other repos. This avoids calling torch.log() on a zero value, which would be undefined. Defaults to 0.03.""" num_images: int = 10000 # TODO: don't hardcode this """How many images exist in the dataset.""" appearance_embedding_dim: int = 48 """Dimension of appearance embedding. Defaults to 48.""" transient_embedding_dim: int = 16 """Dimension of transient embedding. Defaults to 16.""" @dataclass class TensoRFModelConfig(ModelConfig): """TensoRF model config""" _target: Type = TensoRFModel """target class to instantiate""" init_resolution: int = 128 """initial render resolution""" final_resolution: int = 200 """final render resolution""" upsampling_iters: Tuple[int, ...] = (5000, 5500, 7000) """specifies a list of iteration step numbers to perform upsampling""" loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss_coarse": 1.0, "feature_loss": 8e-5}) """Loss specific weights.""" # Pipeline related configs @dataclass class PipelineConfig(InstantiateConfig): """Configuration for pipeline instantiation""" _target: Type = Pipeline """target class to instantiate""" datamanager: VanillaDataManagerConfig = VanillaDataManagerConfig() """specifies the datamanager config""" model: ModelConfig = ModelConfig() """specifies the model config""" # Viewer related configs @dataclass class ViewerConfig(PrintableConfig): """Configuration for viewer instantiation""" log_filename: Optional[Path] = None """Filename to use for the log file. Defaults to None. If None, no log file is created.""" enable: bool = False """whether to enable viewer""" start_train: bool = True """whether to immediately start training upon loading viewer if False, will just visualize dataset but you can toggle training in viewer""" zmq_url: str = "tcp://127.0.0.1:6000" """the zmq port to connect to for communication""" launch_bridge_server: bool = True """whether or not to launch the zmq bridge server""" websocket_port: int = 7007 """the default websocket port to connect to""" num_rays_per_chunk: int = 32768 """number of rays per chunk to render with visualizer""" # Optimizer related configs @dataclass class OptimizerConfig(InstantiateConfig): """Basic optimizer config with RAdam""" _target: Type = torch.optim.RAdam lr: float = 0.0005 eps: float = 1e-08 # TODO: somehow make this more generic. i dont like the idea of overriding the setup function # but also not sure how to go about passing things into predefined torch objects. def setup(self, params=None, **kwargs) -> Any: """Returns the instantiated object using the config.""" return self._target(params, lr=self.lr, eps=self.eps) @dataclass class SchedulerConfig(InstantiateConfig): """Basic scheduler config with self-defined exponential decay schedule""" _target: Type = ExponentialDecaySchedule lr_final: float = 0.000005 max_steps: int = 1000000 # TODO: somehow make this more generic. i dont like the idea of overriding the setup function # but also not sure how to go about passing things into predefined torch objects. def setup(self, optimizer=None, lr_init=None, **kwargs) -> Any: """Returns the instantiated object using the config.""" return self._target(optimizer, lr_init, self.lr_final, self.max_steps) @dataclass class Config(PrintableConfig): """Full config contents""" experiment_name: str = "blender_lego" method_name: str = "base_method" base_dir: Optional[Path] = None # base dir path to be dynamically set machine: MachineConfig = MachineConfig() logging: LoggingConfig = LoggingConfig() trainer: TrainerConfig = TrainerConfig() pipeline: PipelineConfig = PipelineConfig() optimizers: Dict[str, Any] = to_immutable_dict( { "fields": { "optimizer": OptimizerConfig(), "scheduler": SchedulerConfig(), } } ) viewer: ViewerConfig = ViewerConfig() def __post_init__(self): """Convert logging directories to more specific filepaths""" dt_str = datetime.now().strftime("%Y-%m-%d_%H%M%S") self.base_dir = Path(f"outputs/{self.experiment_name}/{self.method_name}/{dt_str}") if self.trainer.model_dir is None: self.trainer.model_dir = self.base_dir / self.trainer.relative_model_dir for curr_writer in self.logging.writer: curr_writer.log_dir = self.base_dir / curr_writer.relative_log_dir if self.viewer.log_filename is None: self.viewer.log_filename = self.base_dir / "viewer_log_filename.txt"