diff --git a/.github/workflows/core_code_checks.yml b/.github/workflows/core_code_checks.yml index 934f4aace..6c4e34608 100644 --- a/.github/workflows/core_code_checks.yml +++ b/.github/workflows/core_code_checks.yml @@ -38,7 +38,7 @@ jobs: run: black docs/ nerfactory/ scripts/ tests/ --check - name: Python Pylint run: | - pylint nerfactory tests scripts + pylint --disable=cyclic-import nerfactory tests scripts - name: Test with pytest run: | pytest diff --git a/.vscode/launch.json b/.vscode/launch.json index 953a075df..a2cef86e5 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -43,7 +43,7 @@ "program": "scripts/run_train.py", "console": "integratedTerminal", "args": [ - "--config-name=graph_instant_ngp.yaml" + "instant_ngp" ] }, { @@ -53,7 +53,7 @@ "program": "scripts/run_train.py", "console": "integratedTerminal", "args": [ - "--config-name=graph_mipnerf_360.yaml" + "mipnerf_360" ] }, { @@ -63,7 +63,7 @@ "program": "scripts/run_train.py", "console": "integratedTerminal", "args": [ - "--config-name=graph_nerfw.yaml" + "nerfw" // "++data.use_preprocessing_cache=true" ] }, @@ -74,7 +74,7 @@ "program": "scripts/run_train.py", "console": "integratedTerminal", "args": [ - "--config-name=graph_semantic_nerf.yaml", + "semantic_nerf", "++data.use_preprocessing_cache=true" ] } diff --git a/README.md b/README.md index 884cb8a53..33f3c0ee1 100644 --- a/README.md +++ b/README.md @@ -156,7 +156,7 @@ logging: #### 5. Rendering a trajectories during inference ``` -python scripts/run_eval.py --method=traj --traj=spiral --output-filename=output.mp4 --config-name=graph_instant_ngp.yaml trainer.resume_train.load_dir=outputs/blender_lego/instant_ngp/2022-07-07_230905/checkpoints +python scripts/run_eval.py --method=traj --traj=spiral --output-filename=output.mp4 --config-name=graph_instant_ngp.yaml --checkpoint-dir=outputs/blender_lego/instant_ngp/2022-07-07_230905/checkpoints ``` #### 6. In-depth guide diff --git a/configs/README.md b/configs/README.md deleted file mode 100644 index e69de29bb..000000000 diff --git a/configs/graph_default.yaml b/configs/graph_default.yaml deleted file mode 100644 index 9b2ddb558..000000000 --- a/configs/graph_default.yaml +++ /dev/null @@ -1,97 +0,0 @@ -defaults: - - override hydra/job_logging: colorlog - - override hydra/hydra_logging: colorlog - - _self_ - -machine: - seed: 42 - num_gpus: 1 - num_machines: 1 - machine_rank: 0 - dist_url: auto - -logging: - steps_per_log: 10 - max_buffer_size: 20 - # writer logs losses and images per iteration - writer: - TensorboardWriter: - log_dir: './' - LocalWriter: - log_dir: './' - stats_to_track: - [ITER_LOAD_TIME, ITER_TRAIN_TIME, RAYS_PER_SEC, CURR_TEST_PSNR] - max_log_size: 10 # if 0, logs everything with no erasing - # profiler logs run times of functions and prints at end of training - enable_profiler: True - -trainer: - model_dir: './checkpoints' - steps_per_save: 1000 - steps_per_test: 500 - max_num_iterations: 1000000 - mixed_precision: False - resume_train: - load_dir: # pretrained model directory - load_step: # pretrained model's step number - -experiment_name: blender_lego -method_name: vanilla_nerf - -# hydra settings -hydra: - mode: RunMode.RUN - run: - dir: outputs/${experiment_name}/${method_name}/${now:%Y-%m-%d_%H%M%S} - job: - chdir: True - -pipeline: - _target_: nerfactory.pipelines.base.Pipeline - dataloader: - _target_: nerfactory.dataloaders.base.VanillaDataloader - image_dataset_type: rgb - train_dataset: - _target_: nerfactory.dataloaders.datasets.Dataset - train_num_rays_per_batch: 1024 - train_num_images_to_sample_from: -1 - eval_image_indices: [0] - eval_num_rays_per_chunk: 4096 - model: - _target_: nerfactory.models.vanilla_nerf.NeRFModel - enable_collider: True - collider_config: - _target_: nerfactory.models.modules.scene_colliders.NearFarCollider - near_plane: 2.0 - far_plane: 6.0 - num_coarse_samples: 64 - num_importance_samples: 128 - loss_coefficients: - rgb_loss_coarse: 1.0 - rgb_loss_fine: 1.0 - enable_density_field: false - density_field_config: - # Please see the docstr for DensityGrid on how to set those arguments. - _target_: nerfactory.fields.density_fields.density_grid.DensityGrid - 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 - -optimizers: - fields: - optimizer: - _target_: torch.optim.RAdam - lr: 0.0005 - scheduler: - _target_: nerfactory.optimizers.optimizers.ExponentialDecaySchedule - lr_final: 0.000005 - max_steps: 1000000 - -viewer: - enable: False - zmq_url: tcp://127.0.0.1:6000 - min_render_image_height: 64 - max_render_image_height: 1024 - num_rays_per_chunk: 4096 diff --git a/configs/graph_instant_ngp.yaml b/configs/graph_instant_ngp.yaml deleted file mode 100644 index 849e8cbe8..000000000 --- a/configs/graph_instant_ngp.yaml +++ /dev/null @@ -1,52 +0,0 @@ -defaults: - - graph_default - - _self_ - -method_name: instant_ngp - -viewer: - enable: true - num_rays_per_chunk: 16384 - -pipeline: - dataloader: - train_dataset: - _target_: nerfactory.dataloaders.datasets.Blender - data_directory: data/blender/lego - alpha_color: white - downscale_factor: 1 - train_num_rays_per_batch: 8192 - eval_num_rays_per_chunk: 8192 - model: - _target_: nerfactory.models.instant_ngp.NGPModel - enable_density_field: true - enable_collider: false - field_implementation: tcnn # torch, tcnn, ... - loss_coefficients: # loss coefficients - rgb_loss: 1.0 - - density_field_config: - # Please see the docstr for DensityGrid on how to set those arguments. - _target_: nerfactory.fields.density_fields.density_grid.DensityGrid - 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 - -# optimizer options for network -optimizers: - fields: - optimizer: - lr: 3e-3 - eps: 1e-15 # NGP prefers small eps - # Note: If you want slightly faster coverange, you can use the following optimizer instead, which - # requires to install apex following https://github.com/NVIDIA/apex#linux - # optimizer: - # _target_: apex.optimizers.FusedAdam - # lr: 3e-3 - # eps: 1e-15 - scheduler: null - -trainer: - mixed_precision: true # Stabilizes training when using tcnn diff --git a/configs/graph_mipnerf.yaml b/configs/graph_mipnerf.yaml deleted file mode 100644 index b8b94c549..000000000 --- a/configs/graph_mipnerf.yaml +++ /dev/null @@ -1,25 +0,0 @@ -defaults: - - graph_default - - _self_ - -method_name: mipnerf - -# Differences compared to paper -# This repo mipNeRF -# density = softplus(x) density = softplus(x-1) -# rgb = sigmoid(x) rgb = (1 + 2e) / (1 + exp(-x)) -e - -pipeline: - dataloader: - train_dataset: - _target_: nerfactory.dataloaders.datasets.Blender - data_directory: data/blender/lego - alpha_color: white - downscale_factor: 1 - model: - _target_: nerfactory.models.mipnerf.MipNerfModel - num_coarse_samples: 128 - num_importance_samples: 128 - loss_coefficients: - rgb_loss_coarse: 0.1 - rgb_loss_fine: 1.0 diff --git a/configs/graph_mipnerf_360.yaml b/configs/graph_mipnerf_360.yaml deleted file mode 100644 index 93d8d431e..000000000 --- a/configs/graph_mipnerf_360.yaml +++ /dev/null @@ -1,27 +0,0 @@ -defaults: - - graph_default - - _self_ - -experiment_name: mipnerf_360 -method_name: vanilla_nerf - -pipeline: - dataloader: - train_dataset: - _target_: nerfactory.dataloaders.datasets.Mipnerf360 - data_directory: data/mipnerf_360/garden - downscale_factor: 1 - model: - _target_: nerfactory.models.mipnerf_360.MipNerf360Model - num_coarse_samples: 128 - num_importance_samples: 128 - collider_config: - _target_: nerfactory.models.modules.scene_colliders.NearFarCollider - near_plane: 0.5 - far_plane: 20.0 - loss_coefficients: - ray_loss_coarse: 1.0 - ray_loss_fine: 1.0 - -trainer: - steps_per_test: 200 diff --git a/configs/graph_nerfw.yaml b/configs/graph_nerfw.yaml deleted file mode 100644 index a36cf4be9..000000000 --- a/configs/graph_nerfw.yaml +++ /dev/null @@ -1,25 +0,0 @@ -defaults: - - graph_default - - _self_ - -experiment_name: friends_TBBT-big_living_room -method_name: nerfw - -pipeline: - dataloader: - image_dataset_type: panoptic - train_dataset: - _target_: nerfactory.dataloaders.datasets.Friends - data_directory: data/friends/TBBT-big_living_room - downscale_factor: 1 - model: - _target_: nerfactory.models.nerfw.NerfWModel - collider_config: - _target_: nerfactory.models.modules.scene_colliders.AABBBoxCollider - num_coarse_samples: 64 - num_importance_samples: 64 - loss_coefficients: - rgb_loss_coarse: 1.0 - rgb_loss_fine: 1.0 - uncertainty_loss: 1.0 - density_loss: 0.01 diff --git a/configs/graph_semantic_nerf.yaml b/configs/graph_semantic_nerf.yaml deleted file mode 100644 index 385aa236b..000000000 --- a/configs/graph_semantic_nerf.yaml +++ /dev/null @@ -1,24 +0,0 @@ -defaults: - - graph_default - - _self_ - -experiment_name: friends_TBBT-big_living_room -method_name: nerfw - -pipeline: - dataloader: - image_dataset_type: panoptic - train_dataset: - _target_: nerfactory.dataloaders.datasets.Friends - data_directory: data/friends/TBBT-big_living_room - downscale_factor: 1 - model: - _target_: nerfactory.models.semantic_nerf.SemanticNerfModel - collider_config: - _target_: nerfactory.models.modules.scene_colliders.AABBBoxCollider - num_coarse_samples: 64 - num_importance_samples: 64 - loss_coefficients: - rgb_loss_coarse: 1.0 - rgb_loss_fine: 1.0 - semantic_loss_fine: 0.05 diff --git a/configs/graph_vanilla_nerf.yaml b/configs/graph_vanilla_nerf.yaml deleted file mode 100644 index cb18b562b..000000000 --- a/configs/graph_vanilla_nerf.yaml +++ /dev/null @@ -1,11 +0,0 @@ -defaults: - - graph_default - - _self_ - -pipeline: - dataloader: - train_dataset: - _target_: nerfactory.dataloaders.datasets.Blender - data_directory: data/blender/lego - alpha_color: white - downscale_factor: 1 diff --git a/docs/conf.py b/docs/conf.py index a6f8562eb..177b8b410 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -15,6 +15,8 @@ import os import sys sys.path.insert(0, os.path.abspath("..")) +import nerfactory.configs.base +import nerfactory.dataloaders.base # -- Project information ----------------------------------------------------- diff --git a/docs/tutorials/quickstart.md b/docs/tutorials/quickstart.md index d67dad81d..16137ca9a 100644 --- a/docs/tutorials/quickstart.md +++ b/docs/tutorials/quickstart.md @@ -118,7 +118,7 @@ bash environment/run_tensorboard.sh To evaluate the trained NeRF, we provide an evaluation script that allows you to do benchmarking (see our [benchmarking workflow](../tooling/benchmarking.md)) or to render out the scene with a custom trajectory and save the output to a video. ```bash -python scripts/run_eval.py --method=traj --traj=spiral --output-filename=output.mp4 --config-name=graph_instant_ngp.yaml trainer.resume_train.load_dir=outputs/blender_lego/instant_ngp/2022-07-07_230905/checkpoints +python scripts/run_eval.py --method=traj --traj=spiral --output-filename=output.mp4 --config-name=graph_instant_ngp.yaml --checkpoint-dir=outputs/blender_lego/instant_ngp/2022-07-07_230905/checkpoints ``` Please note, this quickstart allows you to preform everything in a headless manner. We also provide a web-based visualizer that allows you to easily monitor training or render out trajectories. See our [viewer docs](../tutorials/viewer/viewer_quickstart.md) for more. diff --git a/environment/requirements.txt b/environment/requirements.txt index 3a6806495..80411cbb0 100644 --- a/environment/requirements.txt +++ b/environment/requirements.txt @@ -2,7 +2,7 @@ aiohttp aiortc av black[jupyter]==22.3.0 -dcargs +dcargs>=0.2.5 ninja functorch==0.1.1 h5py>=2.9.0 diff --git a/nerfactory/cameras/rays.py b/nerfactory/cameras/rays.py index 19d119fcb..82771a10e 100644 --- a/nerfactory/cameras/rays.py +++ b/nerfactory/cameras/rays.py @@ -154,7 +154,7 @@ class RayBundle(TensorDataclass): Args: origins: Ray origins (XYZ) directions: Unit ray direction vector - pixel_area: Projected area of pixel a distance 1 away from origin. + pixel_area: Projected area of pixel a distance 1 away from origin camera_indices: Camera indices nears: Distance along ray to start sampling fars: Rays Distance along ray to stop sampling @@ -162,13 +162,14 @@ class RayBundle(TensorDataclass): num_rays_per_chunk: Number of rays per chunk """ - origins: TensorType["num_rays", 3] - directions: TensorType["num_rays", 3] - pixel_area: TensorType["num_rays", 1] - camera_indices: Optional[TensorType["num_rays", 1]] = None - nears: Optional[TensorType["num_rays", 1]] = None - fars: Optional[TensorType["num_rays", 1]] = None - valid_mask: Optional[TensorType["num_rays", 1, bool]] = None + # TODO(ethan): make sure the sizes with ... are correct + origins: TensorType[..., 3] + directions: TensorType[..., 3] + pixel_area: TensorType[..., 1] + camera_indices: Optional[TensorType[..., 1]] = None + nears: Optional[TensorType[..., 1]] = None + fars: Optional[TensorType[..., 1]] = None + valid_mask: Optional[TensorType[..., 1, bool]] = None num_rays_per_chunk: int = 128 def set_camera_indices(self, camera_index: int) -> None: diff --git a/nerfactory/configs/__init__.py b/nerfactory/configs/__init__.py new file mode 100644 index 000000000..02655775e --- /dev/null +++ b/nerfactory/configs/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py new file mode 100644 index 000000000..1ce43232b --- /dev/null +++ b/nerfactory/configs/base.py @@ -0,0 +1,377 @@ +# 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""" + +import enum +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, Literal, Optional, Tuple, Type + +import torch + +from nerfactory.configs.utils import to_immutable_dict +from nerfactory.dataloaders.base import Dataloader, VanillaDataloader +from nerfactory.dataloaders.datasets import Blender, Dataset, Friends, Mipnerf360 +from nerfactory.models.base import Model +from nerfactory.models.instant_ngp import NGPModel +from nerfactory.models.nerfw import NerfWModel +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 + num_gpus: int = 1 + num_machines: int = 1 + machine_rank: int = 0 + dist_url: str = "auto" + + +# Logging related configs +@dataclass +class TensorboardWriterConfig(InstantiateConfig): + """Tensorboard Writer config""" + + _target: Type = writer.TensorboardWriter + enable: bool = False + relative_log_dir: Path = Path("./") + log_dir: Optional[Path] = None # full log dir path to be dynamically set + + +@dataclass +class WandbWriterConfig(InstantiateConfig): + """WandDB Writer config""" + + _target: Type = writer.WandbWriter + enable: bool = False + relative_log_dir: Path = Path("./") + log_dir: Optional[Path] = None # full log dir path to be dynamically set + + +@dataclass +class LocalWriterConfig(InstantiateConfig): + """Local Writer config""" + + _target: Type = writer.LocalWriter + enable: bool = False + stats_to_track: Tuple[writer.EventName, ...] = ( + writer.EventName.ITER_LOAD_TIME, + writer.EventName.ITER_TRAIN_TIME, + writer.EventName.RAYS_PER_SEC, + writer.EventName.CURR_TEST_PSNR, + ) + max_log_size: int = 10 + relative_log_dir: Path = Path("./") + log_dir: Optional[Path] = None # full log dir path to be dynamically set + + +@dataclass +class LoggingConfig(PrintableConfig): + """Configuration of loggers and profilers""" + + steps_per_log: int = 10 + max_buffer_size: int = 20 + writer: Tuple[Any, ...] = ( + TensorboardWriterConfig(enable=True), + WandbWriterConfig(enable=False), + LocalWriterConfig(enable=True), + ) + # profiler logs run times of functions and prints at end of training + enable_profiler: bool = True + + +# Trainer related configs +@dataclass +class TrainerConfig(PrintableConfig): + """Configuration for training regimen""" + + steps_per_save: int = 1000 + steps_per_test: int = 500 + max_num_iterations: int = 1000000 + mixed_precision: bool = False + relative_model_dir: Path = Path("nerfactory_models/") + model_dir: Optional[Path] = None # full model dir path to be dynamically set + # optional parameters if we want to resume training + load_dir: Optional[Path] = None + load_step: Optional[int] = None + load_config: Optional[Path] = None + + +# Dataset related configs +@dataclass +class DatasetConfig(InstantiateConfig): + """Basic dataset config""" + + _target: Type = Dataset + + +@dataclass +class DataloaderConfig(InstantiateConfig): + """Configuration for dataloader instantiation""" + + _target: Type = Dataloader + image_dataset_type: str = "rgb" + train_dataset: InstantiateConfig = DatasetConfig() + train_num_rays_per_batch: int = 1024 + train_num_images_to_sample_from: int = -1 + eval_dataset: Optional[InstantiateConfig] = None + eval_image_indices: Optional[Tuple[int, ...]] = (0,) + eval_num_rays_per_chunk: int = 4096 + + +@dataclass +class BlenderDatasetConfig(InstantiateConfig): + """Blender dataset config""" + + _target: Type = Blender + data_directory: str = "data/blender/lego" + scale_factor: float = 1.0 + alpha_color: str = "white" + downscale_factor: int = 1 + + +@dataclass +class BlenderDataloaderConfig(DataloaderConfig): + """Blender dataloader config""" + + _target: Type = VanillaDataloader + train_dataset: InstantiateConfig = BlenderDatasetConfig() + + +@dataclass +class FriendsDatasetConfig(InstantiateConfig): + """Friends dataset config""" + + _target: Type = Friends + data_directory: str = "data/friends/TBBT-big_living_room" + + +@dataclass +class FriendsDataloaderConfig(DataloaderConfig): + """Friends dataloader config""" + + _target: Type = VanillaDataloader + train_dataset: InstantiateConfig = FriendsDatasetConfig() + image_dataset_type: str = "panoptic" + + +@dataclass +class MipNerf360DatasetConfig(InstantiateConfig): + """Mipnerf 360 dataset config""" + + _target: Type = Mipnerf360 + data_directory: str = "data/mipnerf_360/garden" + + +@dataclass +class MipNerf360DataloaderConfig(DataloaderConfig): + """Mipnerf 360 dataloader config""" + + _target: Type = VanillaDataloader + train_dataset: InstantiateConfig = MipNerf360DatasetConfig() + + +# Model related configs +@dataclass +class ModelConfig(InstantiateConfig): + """Configuration for model instantiation""" + + _target: Type = Model + enable_collider: bool = True + collider_params: Dict[str, float] = to_immutable_dict({"near_plane": 2.0, "far_plane": 6.0}) + loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss_coarse": 1.0, "rgb_loss_fine": 1.0}) + num_coarse_samples: int = 64 + num_importance_samples: int = 128 + field_implementation: Literal["torch", "tcnn"] = "torch" + enable_density_field: bool = False + 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, + } + ) + + +@dataclass +class InstantNGPModelConfig(ModelConfig): + """Instant NGP Model Config""" + + _target: Type = NGPModel + 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}) + + +@dataclass +class NerfWModelConfig(ModelConfig): + """NerfW model config""" + + _target: Type = NerfWModel + 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} + ) + num_coarse_samples: int = 64 + num_importance_samples: int = 64 + uncertainty_min: float = 0.03 + + +# Pipeline related configs +@dataclass +class PipelineConfig(InstantiateConfig): + """Configuration for pipeline instantiation""" + + _target: Type = Pipeline + dataloader: DataloaderConfig = DataloaderConfig() + model: ModelConfig = ModelConfig() + + +# Viewer related configs +@dataclass +class ViewerConfig(PrintableConfig): + """Configuration for viewer instantiation""" + + enable: bool = False + zmq_url: str = "tcp://127.0.0.1:6000" + min_render_image_height: int = 64 + max_render_image_height: int = 1024 + num_rays_per_chunk: int = 4096 + + +# 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) + + +class MethodType(enum.Enum): + """Enum for the method type.""" + + PSNR = enum.auto() + TRAJ = enum.auto() + + +class TrajectoryType(enum.Enum): + """Enum for the trajectory type.""" + + SPIRAL = enum.auto() + INTERP = enum.auto() + + +@dataclass +class EvalConfig(PrintableConfig): + """Boiler plate config for running eval""" + + checkpoint_dir: Optional[Path] = None + rendered_output_name: Optional[str] = None + method: MethodType = MethodType.PSNR + traj: TrajectoryType = TrajectoryType.SPIRAL + output_filename: Path = Path("output.json") + rendered_resolution_scaling_factor: float = 1.0 + load_config: Optional[Path] = None # optionally load config file from training + + +@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() + eval: Optional[EvalConfig] = None + + 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}") + 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 diff --git a/nerfactory/configs/base_configs.py b/nerfactory/configs/base_configs.py new file mode 100644 index 000000000..e5bd22e89 --- /dev/null +++ b/nerfactory/configs/base_configs.py @@ -0,0 +1,104 @@ +# 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. + +""" +Put all the method implementations in one location. +""" + +from nerfactory.configs.base import ( + BlenderDataloaderConfig, + Config, + FriendsDataloaderConfig, + InstantNGPModelConfig, + MipNerf360DataloaderConfig, + ModelConfig, + NerfWModelConfig, + OptimizerConfig, + PipelineConfig, + TrainerConfig, +) +from nerfactory.models.mipnerf import MipNerfModel +from nerfactory.models.mipnerf_360 import MipNerf360Model +from nerfactory.models.semantic_nerf import SemanticNerfModel +from nerfactory.models.vanilla_nerf import NeRFModel + +base_configs = {} +base_configs["instant_ngp"] = Config( + method_name="instant_ngp", + trainer=TrainerConfig(mixed_precision=True), + pipeline=PipelineConfig( + dataloader=BlenderDataloaderConfig(train_num_rays_per_batch=8192, eval_num_rays_per_chunk=8192), + model=InstantNGPModelConfig(), + ), + optimizers={ + "fields": { + "optimizer": OptimizerConfig(lr=3e-3, eps=1e-15), + "scheduler": None, + } + }, +) + +base_configs["mipnerf_360"] = Config( + experiment_name="mipnerf_360", + method_name="mipnerf_360", + trainer=TrainerConfig(steps_per_test=200), + pipeline=PipelineConfig( + dataloader=MipNerf360DataloaderConfig(), + model=ModelConfig( + _target=MipNerf360Model, + collider_params={"near_plane": 0.5, "far_plane": 20.0}, + loss_coefficients={"ray_loss_coarse": 1.0, "ray_loss_fine": 1.0}, + num_coarse_samples=128, + num_importance_samples=128, + ), + ), +) + +base_configs["mipnerf"] = Config( + method_name="mipnerf", + pipeline=PipelineConfig( + dataloader=BlenderDataloaderConfig(), + model=ModelConfig( + _target=MipNerfModel, + loss_coefficients={"rgb_loss_coarse": 0.1, "rgb_loss_fine": 1.0}, + num_coarse_samples=128, + num_importance_samples=128, + ), + ), +) + +base_configs["nerfw"] = Config( + experiment_name="friends_TBBT-big_living_room", + method_name="nerfw", + pipeline=PipelineConfig(dataloader=FriendsDataloaderConfig(), model=NerfWModelConfig()), +) + +base_configs["semantic_nerf"] = Config( + experiment_name="friends_TBBT-big_living_room", + method_name="semantic_nerf", + pipeline=PipelineConfig( + dataloader=FriendsDataloaderConfig(), + model=ModelConfig( + _target=SemanticNerfModel, + loss_coefficients={"rgb_loss_coarse": 1.0, "rgb_loss_fine": 1.0, "semantic_loss_fine": 0.05}, + num_coarse_samples=64, + num_importance_samples=64, + ), + ), +) + +base_configs["vanilla_nerf"] = Config( + method_name="vanilla_nerf", + pipeline=PipelineConfig(dataloader=BlenderDataloaderConfig(), model=ModelConfig(_target=NeRFModel)), +) diff --git a/nerfactory/configs/utils.py b/nerfactory/configs/utils.py new file mode 100644 index 000000000..e2c1567d3 --- /dev/null +++ b/nerfactory/configs/utils.py @@ -0,0 +1,70 @@ +# 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. + +""" +Some utility code for configs. +""" + +from __future__ import annotations + +import sys +from dataclasses import field +from typing import Any, Dict, TypeVar + +import dcargs + +from nerfactory.configs import base as cfg + +T = TypeVar("T") + + +def cli_from_base_configs(base_library: Dict[str, T], eval_mode=False) -> T: + """Populate an instance of `cls`, where the first positional argument is used to + select from a library of named base configs. See + https://brentyi.github.io/dcargs/examples/06_base_configs/ for more details.""" + + if eval_mode: + for _, v in base_library.items(): + v.eval = cfg.EvalConfig() + + # Get base configuration name from the first positional argument. + if len(sys.argv) < 2 or sys.argv[1] not in base_library: + valid_usages = map(lambda k: f"{sys.argv[0]} {k} --help", base_library.keys()) + raise SystemExit("usage:\n " + "\n ".join(valid_usages)) + + # Get base configuration from our library, and use it for default CLI parameters. + default_instance = base_library[sys.argv[1]] + + return dcargs.cli( + type(default_instance), + prog=" ".join(sys.argv[:2]), + args=sys.argv[2:], + default_instance=default_instance, + # `avoid_subparsers` will avoid making a subparser for unions when a default is + # provided; in this case, it simplifies our CLI but makes it less expressive + # (cannot switch away from the base optimizer types). + avoid_subparsers=True, + ) + + +# pylint: disable=import-outside-toplevel + +# cannot use mutable types directly within dataclass; abstracting default factory calls +def to_immutable_dict(d: Dict[str, Any]): + """Method to convert mutable dict to default factory dict + + Args: + d: dictionary to convert into default factory dict for dataclass + """ + return field(default_factory=lambda: dict(d)) diff --git a/nerfactory/dataloaders/base.py b/nerfactory/dataloaders/base.py index 46a3f87ce..8302cbd70 100644 --- a/nerfactory/dataloaders/base.py +++ b/nerfactory/dataloaders/base.py @@ -15,26 +15,24 @@ """ Data loader. """ +from __future__ import annotations import logging from abc import abstractmethod from typing import Dict, List, Tuple, Union import torch -from omegaconf import DictConfig, ListConfig from torch import nn from torch.nn import Parameter from nerfactory.cameras.rays import RayBundle +from nerfactory.configs import base as cfg from nerfactory.dataloaders.eval import FixedIndicesEvalDataloader from nerfactory.dataloaders.image_dataset import ImageDataset, PanopticImageDataset from nerfactory.dataloaders.image_sampler import CacheImageSampler from nerfactory.dataloaders.pixel_sampler import PixelSampler -from nerfactory.dataloaders.structs import DatasetInputs from nerfactory.models.modules.ray_generator import RayGenerator -from nerfactory.utils import profiler -from nerfactory.utils.config import DataloaderConfig -from nerfactory.utils.misc import IterableWrapper, instantiate_from_dict_config +from nerfactory.utils.misc import IterableWrapper class Dataloader(nn.Module): @@ -198,57 +196,59 @@ class VanillaDataloader(Dataloader): # pylint: disable=abstract-method only the constructor is likely to change in the future, or maybe passing in step number to the next_train and next_eval functions.""" + config: cfg.DataloaderConfig + def __init__( self, - image_dataset_type: str, - train_datasetinputs: DatasetInputs, - train_num_rays_per_batch: int, - train_num_images_to_sample_from: int, - eval_datasetinputs: DatasetInputs, - eval_image_indices: Union[List[int], ListConfig], # TODO(ethan): get rid of this hydra ListConfig nonsense - eval_num_rays_per_chunk: int, + config: cfg.DataloaderConfig, device: Union[torch.device, str] = "cpu", - **kwargs # pylint: disable=unused-argument + test_mode: bool = False, + **kwargs, # pylint: disable=unused-argument ): - self.image_dataset_type = image_dataset_type - self.train_datasetinputs = train_datasetinputs - self.train_num_rays_per_batch = train_num_rays_per_batch - self.train_num_images_to_sample_from = train_num_images_to_sample_from - self.eval_datasetinputs = eval_datasetinputs - self.eval_image_indices = eval_image_indices - self.eval_num_rays_per_chunk = eval_num_rays_per_chunk + self.config = config self.device = device + + dataset_train = config.train_dataset.setup() + self.train_datasetinputs = dataset_train.get_dataset_inputs(split="train") + if config.eval_dataset is not None: + dataset_eval = config.eval_dataset.setup() + else: + logging.info("No eval dataset specified so using train dataset for eval.") + dataset_eval = dataset_train + self.config.eval_datasetinputs = dataset_eval.get_dataset_inputs(split="val" if not test_mode else "test") use_train = self.train_datasetinputs is not None - use_eval = self.eval_datasetinputs is not None + use_eval = self.config.eval_datasetinputs is not None super().__init__(use_train, use_eval) def setup_train(self): """Sets up the dataloader for training""" - if self.image_dataset_type == "rgb": + if self.config.image_dataset_type == "rgb": self.train_image_dataset = ImageDataset(**self.train_datasetinputs.as_dict()) - elif self.image_dataset_type == "panoptic": + elif self.config.image_dataset_type == "panoptic": self.train_image_dataset = PanopticImageDataset(**self.train_datasetinputs.as_dict()) self.train_image_sampler = CacheImageSampler( - self.train_image_dataset, num_images_to_sample_from=self.train_num_images_to_sample_from, device=self.device + self.train_image_dataset, + num_images_to_sample_from=self.config.train_num_images_to_sample_from, + device=self.device, ) # TODO(ethan): pass this in self.iter_train_image_sampler = iter(self.train_image_sampler) - self.train_pixel_sampler = PixelSampler(self.train_num_rays_per_batch) + self.train_pixel_sampler = PixelSampler(self.config.train_num_rays_per_batch) self.train_ray_generator = RayGenerator( self.train_datasetinputs.intrinsics, self.train_datasetinputs.camera_to_world ) def setup_eval(self): """Sets up the dataloader for evaluation""" - if self.image_dataset_type == "rgb": - self.eval_image_dataset = ImageDataset(**self.eval_datasetinputs.as_dict()) - elif self.image_dataset_type == "panoptic": - self.eval_image_dataset = PanopticImageDataset(**self.eval_datasetinputs.as_dict()) + if self.config.image_dataset_type == "rgb": + self.eval_image_dataset = ImageDataset(**self.config.eval_datasetinputs.as_dict()) + elif self.config.image_dataset_type == "panoptic": + self.eval_image_dataset = PanopticImageDataset(**self.config.eval_datasetinputs.as_dict()) self.eval_dataloader = FixedIndicesEvalDataloader( image_dataset=self.eval_image_dataset, - intrinsics=self.eval_datasetinputs.intrinsics, - camera_to_world=self.eval_datasetinputs.camera_to_world, - num_rays_per_chunk=self.eval_num_rays_per_chunk, - image_indices=self.eval_image_indices, + intrinsics=self.config.eval_datasetinputs.intrinsics, + camera_to_world=self.config.eval_datasetinputs.camera_to_world, + num_rays_per_chunk=self.config.eval_num_rays_per_chunk, + image_indices=self.config.eval_image_indices, device=self.device, ) @@ -268,25 +268,3 @@ class VanillaDataloader(Dataloader): # pylint: disable=abstract-method self.eval_count += 1 camera_ray_bundle, batch = next(self.eval_dataloader) return camera_ray_bundle, batch - - -@profiler.time_function -def setup_dataloader(config: DataloaderConfig, device: str, test_mode=False) -> Dataloader: - """Setup the dataloader.""" - dataset_train = instantiate_from_dict_config(config.train_dataset) - train_datasetinputs = dataset_train.get_dataset_inputs(split="train") - if config.eval_dataset is not None: - dataset_eval = instantiate_from_dict_config(config.eval_dataset) - else: - logging.info("No eval dataset specified so using train dataset for eval.") - dataset_eval = dataset_train - eval_datasetinputs = dataset_eval.get_dataset_inputs(split="val" if not test_mode else "test") - - dataloader: Dataloader = instantiate_from_dict_config( - DictConfig(config), - train_datasetinputs=train_datasetinputs, - eval_datasetinputs=eval_datasetinputs, - device=device, - ) - dataloader.to(device) - return dataloader diff --git a/nerfactory/dataloaders/datasets.py b/nerfactory/dataloaders/datasets.py index 63425ab11..283635723 100644 --- a/nerfactory/dataloaders/datasets.py +++ b/nerfactory/dataloaders/datasets.py @@ -14,18 +14,20 @@ """A set of standard datasets.""" +from __future__ import annotations + import dataclasses import logging import os from abc import abstractmethod from dataclasses import dataclass -from typing import Optional, Union import imageio import numpy as np import torch from scipy.spatial.transform import Rotation +from nerfactory.configs import base as cfg from nerfactory.dataloaders.colmap_utils import ( read_cameras_binary, read_images_binary, @@ -133,10 +135,13 @@ class Blender(Dataset): downscale_factor: How much to downscale images. Defaults to 1. """ - data_directory: str - scale_factor: float = 1.0 - alpha_color: Optional[Union[str, list]] = None - downscale_factor: int = 1 + def __init__(self, config: cfg.BlenderDatasetConfig): + super().__init__() + self.config = config + self.data_directory: str = config.data_directory + self.scale_factor: float = config.scale_factor + self.alpha_color = config.alpha_color + self.downscale_factor: int = config.downscale_factor def _generate_dataset_inputs(self, split="train"): if self.alpha_color is not None: diff --git a/nerfactory/dataloaders/eval.py b/nerfactory/dataloaders/eval.py index c55284e22..5e589686c 100644 --- a/nerfactory/dataloaders/eval.py +++ b/nerfactory/dataloaders/eval.py @@ -18,10 +18,9 @@ Data loader. import random from abc import abstractmethod -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, Optional, Tuple, Union import torch -from omegaconf import ListConfig from nerfactory.cameras.cameras import Camera, get_camera from nerfactory.cameras.rays import RayBundle @@ -84,7 +83,7 @@ class FixedIndicesEvalDataloader(EvalDataloader): intrinsics, camera_to_world, num_rays_per_chunk: int, - image_indices: Optional[Union[List[int], ListConfig]] = None, + image_indices: Optional[Tuple[int]] = None, device: Union[torch.device, str] = "cpu", **kwargs, ): diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index 6795327ca..fe79cf13c 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -15,10 +15,13 @@ """ Code to train model. """ +from __future__ import annotations + import functools import logging import os import typing +from pathlib import Path from typing import Any, Dict, List import torch @@ -26,12 +29,12 @@ import torch.distributed as dist from torch.cuda.amp.grad_scaler import GradScaler from torch.nn.parallel import DistributedDataParallel as DDP +from nerfactory.configs import base as cfg from nerfactory.dataloaders.structs import DatasetInputs from nerfactory.optimizers.optimizers import Optimizers, setup_optimizers -from nerfactory.pipelines.base import Pipeline, setup_pipeline +from nerfactory.pipelines.base import Pipeline from nerfactory.utils import profiler, writer from nerfactory.utils.callbacks import Callback -from nerfactory.utils.config import Config from nerfactory.utils.decorators import check_main_thread from nerfactory.utils.writer import EventName, TimeWriter from nerfactory.viewer.server import viewer_utils @@ -39,7 +42,7 @@ from nerfactory.viewer.server import viewer_utils logging.getLogger("PIL").setLevel(logging.WARNING) -def train_loop(local_rank: int, world_size: int, config: Config) -> Any: +def train_loop(local_rank: int, world_size: int, config: cfg.Config) -> Any: """Main training function that sets up and runs the trainer per process Args: @@ -65,7 +68,7 @@ class Trainer: world_size (int, optional): World size of the process. Defaults to 1. """ - def __init__(self, config: Config, local_rank: int = 0, world_size: int = 1): + def __init__(self, config: cfg.Config, local_rank: int = 0, world_size: int = 1): self.config = config self.local_rank = local_rank self.world_size = world_size @@ -95,7 +98,7 @@ class Trainer: Args: test_mode (bool, optional): Whether to setup for testing. Defaults to False. """ - self.pipeline: Pipeline = setup_pipeline(self.config.pipeline, device=self.device, test_mode=test_mode) + self.pipeline: Pipeline = self.config.pipeline.setup(device=self.device, test_mode=test_mode) self.optimizers = setup_optimizers(self.config.optimizers, self.pipeline.get_param_groups()) self._load_checkpoint() @@ -125,7 +128,7 @@ class Trainer: with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as t: self.visualizer_state.update_scene(step, self.pipeline.model) - train_num_rays_per_batch = self.pipeline.dataloader.train_num_rays_per_batch + train_num_rays_per_batch = self.config.pipeline.dataloader.train_num_rays_per_batch writer.put_scalar(name=EventName.RAYS_PER_SEC, scalar=train_num_rays_per_batch / t.duration, step=step) if step != 0 and step % self.config.logging.steps_per_log == 0: @@ -154,9 +157,10 @@ class Trainer: def _load_checkpoint(self) -> None: """Helper function to load pipeline and optimizer from prespecified checkpoint""" - load_config = self.config.trainer.resume_train - if load_config.load_dir is not None and load_config.load_step is not None: - load_path = os.path.join(load_config.load_dir, f"step-{load_config.load_step:09d}.ckpt") + load_dir = self.config.trainer.load_dir + load_step = self.config.trainer.load_step + if load_dir is not None and load_step is not None: + load_path = os.path.join(load_dir, f"step-{load_step:09d}.ckpt") assert os.path.exists(load_path), f"Checkpoint {load_path} does not exist" loaded_state = torch.load(load_path, map_location="cpu") self.start_step = loaded_state["step"] + 1 @@ -169,16 +173,16 @@ class Trainer: logging.info("No checkpoints to load, training from scratch") @check_main_thread - def _save_checkpoint(self, output_dir: str, step: int) -> None: + def _save_checkpoint(self, output_dir: Path, step: int) -> None: """Save the model and optimizers Args: output_dir: directory to save the checkpoint step: number of steps in training for given checkpoint """ - if not os.path.exists(output_dir): - os.makedirs(output_dir) - ckpt_path = os.path.join(output_dir, f"step-{step:09d}.ckpt") + if not output_dir.exists(): + output_dir.mkdir(parents=True, exist_ok=True) + ckpt_path = output_dir / f"step-{step:09d}.ckpt" if hasattr(self.pipeline, "module"): pipeline = self.pipeline.module.state_dict() # type: ignore else: diff --git a/nerfactory/fields/density_fields/density_grid.py b/nerfactory/fields/density_fields/density_grid.py index 9ead64cb3..3a4ee3a8e 100644 --- a/nerfactory/fields/density_fields/density_grid.py +++ b/nerfactory/fields/density_fields/density_grid.py @@ -15,6 +15,9 @@ """ Code to implement the density grid. """ + +from __future__ import annotations + from typing import Callable, List, Tuple, Union import torch @@ -81,12 +84,7 @@ class DensityGrid(nn.Module): mean_density: TensorType[1] def __init__( - self, - center: float = 0.0, - base_scale: float = 1.0, - num_cascades: int = 1, - resolution: int = 128, - update_every_num_iters: int = 16, + self, center: float, base_scale: float, num_cascades: int, resolution: int, update_every_num_iters: int ) -> None: super().__init__() self.center = center @@ -102,7 +100,7 @@ class DensityGrid(nn.Module): self.register_buffer("density_bitfield", density_bitfield) # Integer grid coords / indices that do not related to cascades - grid_coords = _create_grid_coords(resolution) + grid_coords = _create_grid_coords(self.resolution) # TODO(ruilongli): hacky way for now. support cpu version of morton3D? grid_indices = nerfactory_cuda.morton3D(grid_coords.to("cuda:0")).to("cpu") self.register_buffer("grid_coords", grid_coords) diff --git a/nerfactory/models/base.py b/nerfactory/models/base.py index 29be1eee0..9ba18ea17 100644 --- a/nerfactory/models/base.py +++ b/nerfactory/models/base.py @@ -15,26 +15,24 @@ """ Base Model implementation which takes in RayBundles """ + +from __future__ import annotations + from abc import abstractmethod from collections import defaultdict from typing import Any, Dict, List, Optional, Tuple, Union, overload import torch -from omegaconf import DictConfig from torch import nn from torch.nn import Parameter from nerfactory.cameras.cameras import Camera from nerfactory.cameras.rays import RayBundle +from nerfactory.configs import base as cfg from nerfactory.dataloaders.structs import SceneBounds -from nerfactory.utils import profiler +from nerfactory.fields.density_fields.density_grid import DensityGrid from nerfactory.utils.callbacks import Callback -from nerfactory.utils.config import ModelConfig -from nerfactory.utils.misc import ( - get_masked_dict, - instantiate_from_dict_config, - is_not_none, -) +from nerfactory.utils.misc import get_masked_dict, is_not_none class Model(nn.Module): @@ -53,28 +51,21 @@ class Model(nn.Module): density_field_config: Configuration of density field. """ + config: cfg.ModelConfig + def __init__( self, + config: cfg.ModelConfig, scene_bounds: SceneBounds, - loss_coefficients: DictConfig = DictConfig({}), - enable_collider: bool = True, - collider_config: Optional[DictConfig] = None, - enable_density_field: bool = False, - density_field_config: Optional[DictConfig] = None, **kwargs, ) -> None: super().__init__() + self.config = config self.scene_bounds = scene_bounds - self.loss_coefficients = loss_coefficients - self.enable_collider = enable_collider - self.collider_config = collider_config - self.enable_density_field = enable_density_field - self.density_field_config = density_field_config self.density_field = None self.kwargs = kwargs self.collider = None self.populate_density_field() - self.populate_collider() self.populate_fields() self.populate_misc_modules() # populate the modules self.callbacks = None @@ -92,13 +83,15 @@ class Model(nn.Module): def populate_density_field(self): """Set the scene density field to use.""" - if self.enable_density_field: - self.density_field = instantiate_from_dict_config(self.density_field_config) + if self.config.enable_density_field: - def populate_collider(self): - """Set the scene bounds collider to use.""" - if self.enable_collider: - self.collider = instantiate_from_dict_config(self.collider_config, scene_bounds=self.scene_bounds) + self.density_field = DensityGrid( + center=self.config.density_field_params["center"], + base_scale=self.config.density_field_params["base_scale"], + num_cascades=self.config.density_field_params["num_cascades"], + resolution=self.config.density_field_params["resolution"], + update_every_num_iters=self.config.density_field_params["update_every_num_iters"], + ) @abstractmethod def populate_fields(self): @@ -152,7 +145,7 @@ class Model(nn.Module): This outputs different things depending on the configuration of the model and whether or not the batch is provided (whether or not we are training basically).""" if self.collider is not None: - intersected_ray_bundle = self.collider(ray_bundle) + intersected_ray_bundle = self.collider(ray_bundle) # pylint: disable=not-callable valid_mask = intersected_ray_bundle.valid_mask[..., 0] else: # NOTE(ruilongli): we don't need collider for ngp @@ -172,9 +165,8 @@ class Model(nn.Module): outputs = self.get_outputs(intersected_ray_bundle) metrics_dict = self.get_metrics_dict(outputs=outputs, batch=batch) loss_dict = self.get_loss_dict( - outputs=outputs, batch=batch, metrics_dict=metrics_dict, loss_coefficients=self.loss_coefficients + outputs=outputs, batch=batch, metrics_dict=metrics_dict, loss_coefficients=self.config.loss_coefficients ) - return outputs, loss_dict, metrics_dict def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]: @@ -230,15 +222,3 @@ class Model(nn.Module): """Load the checkpoint from the given path""" state = {key.replace("module.", ""): value for key, value in loaded_state["model"].items()} self.load_state_dict(state) # type: ignore - - -@profiler.time_function -def setup_model(config: ModelConfig, scene_bounds: SceneBounds, device: str) -> Model: - """Setup the model. The dataset inputs should be set with the training data. - - Args: - dataset_inputs: The inputs which will be used to define the camera parameters. - """ - model = instantiate_from_dict_config(DictConfig(config), scene_bounds=scene_bounds, device=device) - model.to(device) - return model diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index 36a4c6f66..c12823d4d 100644 --- a/nerfactory/models/instant_ngp.py +++ b/nerfactory/models/instant_ngp.py @@ -16,6 +16,8 @@ Implementation of Instant NGP. """ +from __future__ import annotations + from typing import Dict, List import torch @@ -26,6 +28,7 @@ from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity import nerfactory.cuda as nerfactory_cuda from nerfactory.cameras.rays import RayBundle +from nerfactory.configs import base as cfg from nerfactory.fields.instant_ngp_field import field_implementation_to_class from nerfactory.fields.modules.field_heads import FieldHeadNames from nerfactory.models.base import Model @@ -43,11 +46,10 @@ class NGPModel(Model): kwargs: additional params to pass up to the parent class model """ - def __init__(self, field_implementation="torch", **kwargs) -> None: - assert field_implementation in field_implementation_to_class - self.field_implementation = field_implementation + def __init__(self, config: cfg.InstantNGPModelConfig, **kwargs) -> None: + assert config.field_implementation in field_implementation_to_class self.field = None - super().__init__(**kwargs) + super().__init__(config=config, **kwargs) def get_training_callbacks(self) -> List[Callback]: assert self.density_field is not None @@ -62,7 +64,7 @@ class NGPModel(Model): def populate_fields(self): """Set the fields.""" # torch or tiny-cuda-nn version - self.field = field_implementation_to_class[self.field_implementation](self.scene_bounds.aabb) + self.field = field_implementation_to_class[self.config.field_implementation](self.scene_bounds.aabb) def populate_misc_modules(self): # samplers @@ -76,6 +78,8 @@ class NGPModel(Model): self.ssim = structural_similarity_index_measure self.lpips = LearnedPerceptualImagePatchSimilarity() + # no colliders default + def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} if self.field is None: diff --git a/nerfactory/models/mipnerf.py b/nerfactory/models/mipnerf.py index d28e293f1..dc4bba8d6 100644 --- a/nerfactory/models/mipnerf.py +++ b/nerfactory/models/mipnerf.py @@ -15,7 +15,7 @@ """ Implementation of mip-NeRF. """ - +from __future__ import annotations from typing import Dict, List @@ -26,11 +26,13 @@ from torchmetrics.functional import structural_similarity_index_measure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from nerfactory.cameras.rays import RayBundle +from nerfactory.configs import base as cfg from nerfactory.fields.modules.encoding import NeRFEncoding from nerfactory.fields.modules.field_heads import FieldHeadNames from nerfactory.fields.nerf_field import NeRFField from nerfactory.models.base import Model from nerfactory.models.modules.ray_sampler import PDFSampler, UniformSampler +from nerfactory.models.modules.scene_colliders import NearFarCollider from nerfactory.optimizers.loss import MSELoss from nerfactory.renderers.renderers import ( AccumulationRenderer, @@ -45,18 +47,11 @@ class MipNerfModel(Model): def __init__( self, - near_plane=2.0, - far_plane=6.0, - num_coarse_samples=64, - num_importance_samples=128, + config: cfg.ModelConfig, **kwargs, ) -> None: - self.near_plane = near_plane - self.far_plane = far_plane - self.num_coarse_samples = num_coarse_samples - self.num_importance_samples = num_importance_samples self.field = None - super().__init__(**kwargs) + super().__init__(config=config, **kwargs) def populate_fields(self): """Set the fields.""" @@ -74,8 +69,8 @@ class MipNerfModel(Model): def populate_misc_modules(self): # samplers - self.sampler_uniform = UniformSampler(num_samples=self.num_coarse_samples) - self.sampler_pdf = PDFSampler(num_samples=self.num_importance_samples, include_original=False) + self.sampler_uniform = UniformSampler(num_samples=self.config.num_coarse_samples) + self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples, include_original=False) # renderers self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) @@ -90,6 +85,12 @@ class MipNerfModel(Model): self.ssim = structural_similarity_index_measure self.lpips = LearnedPerceptualImagePatchSimilarity() + # colliders + if self.config.enable_collider: + self.collider = NearFarCollider( + near_plane=self.config.collider_params["near_plane"], far_plane=self.config.collider_params["far_plane"] + ) + def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} if self.field is None: @@ -155,14 +156,14 @@ class MipNerfModel(Model): depth_coarse = visualization.apply_depth_colormap( outputs["depth_coarse"], accumulation=outputs["accumulation_coarse"], - near_plane=self.near_plane, - far_plane=self.far_plane, + near_plane=self.config.collider_params["near_plane"], + far_plane=self.config.collider_params["far_plane"], ) depth_fine = visualization.apply_depth_colormap( outputs["depth_fine"], accumulation=outputs["accumulation_fine"], - near_plane=self.near_plane, - far_plane=self.far_plane, + near_plane=self.config.collider_params["near_plane"], + far_plane=self.config.collider_params["far_plane"], ) combined_rgb = torch.cat([image, rgb_coarse, rgb_fine], dim=1) diff --git a/nerfactory/models/mipnerf_360.py b/nerfactory/models/mipnerf_360.py index 41447d999..8ea2cd086 100644 --- a/nerfactory/models/mipnerf_360.py +++ b/nerfactory/models/mipnerf_360.py @@ -16,6 +16,7 @@ Implementation of mip-NeRF. """ +from __future__ import annotations from typing import Dict, List @@ -26,6 +27,7 @@ from torchmetrics.functional import structural_similarity_index_measure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from nerfactory.cameras.rays import RayBundle +from nerfactory.configs import base as cfg from nerfactory.fields.modules.encoding import NeRFEncoding from nerfactory.fields.modules.field_heads import FieldHeadNames from nerfactory.fields.modules.spatial_distortions import SceneContraction @@ -33,6 +35,7 @@ from nerfactory.fields.nerf_field import NeRFField from nerfactory.models.base import Model from nerfactory.models.modules.ray_losses import distortion_loss from nerfactory.models.modules.ray_sampler import PDFSampler, UniformSampler +from nerfactory.models.modules.scene_colliders import NearFarCollider from nerfactory.optimizers.loss import MSELoss from nerfactory.renderers.renderers import ( AccumulationRenderer, @@ -47,18 +50,11 @@ class MipNerf360Model(Model): def __init__( self, - near_plane=2.0, - far_plane=6.0, - num_coarse_samples=64, - num_importance_samples=128, + config: cfg.ModelConfig, **kwargs, ) -> None: - self.near_plane = near_plane - self.far_plane = far_plane - self.num_coarse_samples = num_coarse_samples - self.num_importance_samples = num_importance_samples self.field = None - super().__init__(**kwargs) + super().__init__(config=config, **kwargs) def populate_fields(self): """Set the fields.""" @@ -79,8 +75,8 @@ class MipNerf360Model(Model): def populate_misc_modules(self): # samplers - self.sampler_uniform = UniformSampler(num_samples=self.num_coarse_samples) - self.sampler_pdf = PDFSampler(num_samples=self.num_importance_samples, include_original=False) + self.sampler_uniform = UniformSampler(num_samples=self.config.num_coarse_samples) + self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples, include_original=False) # renderers self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) @@ -95,6 +91,12 @@ class MipNerf360Model(Model): self.ssim = structural_similarity_index_measure self.lpips = LearnedPerceptualImagePatchSimilarity() + # colliders + if self.config.enable_collider: + self.collider = NearFarCollider( + near_plane=self.config.collider_params["near_plane"], far_plane=self.config.collider_params["far_plane"] + ) + def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} if self.field is None: @@ -169,14 +171,14 @@ class MipNerf360Model(Model): depth_coarse = visualization.apply_depth_colormap( outputs["depth_coarse"], accumulation=outputs["accumulation_coarse"], - near_plane=self.near_plane, - far_plane=self.far_plane, + near_plane=self.config.collider_params["near_plane"], + far_plane=self.config.collider_params["far_plane"], ) depth_fine = visualization.apply_depth_colormap( outputs["depth_fine"], accumulation=outputs["accumulation_fine"], - near_plane=self.near_plane, - far_plane=self.far_plane, + near_plane=self.config.collider_params["near_plane"], + far_plane=self.config.collider_params["far_plane"], ) combined_rgb = torch.cat([image, rgb_coarse, rgb_fine], dim=1) diff --git a/nerfactory/models/modules/scene_colliders.py b/nerfactory/models/modules/scene_colliders.py index 146cf5bb0..df3f550fe 100644 --- a/nerfactory/models/modules/scene_colliders.py +++ b/nerfactory/models/modules/scene_colliders.py @@ -15,6 +15,9 @@ """ Ray generator. """ + +from __future__ import annotations + import torch from torch import nn from torchtyping import TensorType @@ -96,7 +99,7 @@ class AABBBoxCollider(SceneBoundsCollider): class NearFarCollider(SceneBoundsCollider): """Sets the nears and fars with fixed values.""" - def __init__(self, near_plane, far_plane, **kwargs) -> None: + def __init__(self, near_plane: float, far_plane: float, **kwargs) -> None: self.near_plane = near_plane self.far_plane = far_plane super().__init__(**kwargs) diff --git a/nerfactory/models/nerfw.py b/nerfactory/models/nerfw.py index 6466031a1..4017f4022 100644 --- a/nerfactory/models/nerfw.py +++ b/nerfactory/models/nerfw.py @@ -16,18 +16,22 @@ NeRF-W (NeRF in the wild) implementation. """ +from __future__ import annotations + import torch from torchmetrics import PeakSignalNoiseRatio from torchmetrics.functional import structural_similarity_index_measure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from nerfactory.cameras.rays import RayBundle +from nerfactory.configs import base as cfg from nerfactory.fields.modules.encoding import NeRFEncoding from nerfactory.fields.modules.field_heads import FieldHeadNames from nerfactory.fields.nerf_field import NeRFField from nerfactory.fields.nerfw_field import VanillaNerfWField from nerfactory.models.base import Model from nerfactory.models.modules.ray_sampler import PDFSampler, UniformSampler +from nerfactory.models.modules.scene_colliders import AABBBoxCollider from nerfactory.optimizers.loss import MSELoss from nerfactory.renderers.renderers import ( AccumulationRenderer, @@ -43,11 +47,7 @@ class NerfWModel(Model): def __init__( self, - near_plane=2.0, - far_plane=6.0, - num_coarse_samples=64, - num_importance_samples=64, - uncertainty_min=0.03, + config: cfg.NerfWModelConfig, **kwargs, ) -> None: """A NeRF-W model. @@ -59,17 +59,12 @@ class NerfWModel(Model): This avoids calling torch.log() on a zero value, which would be undefined. Defaults to 0.03. """ - self.near_plane = near_plane - self.far_plane = far_plane - self.num_coarse_samples = num_coarse_samples - self.num_importance_samples = num_importance_samples - self.uncertainty_min = uncertainty_min self.field_coarse = None self.field_fine = None self.num_images = 10000 # TODO(ethan): fix this self.appearance_embedding_dim = 48 self.transient_embedding_dim = 16 - super().__init__(**kwargs) + super().__init__(config=config, **kwargs) def populate_fields(self): """Set the fields.""" @@ -92,8 +87,8 @@ class NerfWModel(Model): def populate_misc_modules(self): # samplers - self.sampler_uniform = UniformSampler(num_samples=self.num_coarse_samples) - self.sampler_pdf = PDFSampler(num_samples=self.num_importance_samples) + self.sampler_uniform = UniformSampler(num_samples=self.config.num_coarse_samples) + self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples) # renderers self.renderer_rgb = RGBRenderer(background_color=colors.BLACK) @@ -109,6 +104,10 @@ class NerfWModel(Model): self.ssim = structural_similarity_index_measure self.lpips = LearnedPerceptualImagePatchSimilarity() + # colliders + if self.config.enable_collider: + self.collider = AABBBoxCollider(scene_bounds=self.scene_bounds) + def get_param_groups(self): param_groups = {} if self.field_coarse is None or self.field_fine is None: @@ -176,7 +175,7 @@ class NerfWModel(Model): # uncertainty uncertainty = self.renderer_uncertainty(field_outputs_fine[FieldHeadNames.UNCERTAINTY], weights_fine_transient) - uncertainty += self.uncertainty_min + uncertainty += self.config.uncertainty_min outputs = { "rgb_coarse": rgb_coarse, # (num_rays, 3) diff --git a/nerfactory/models/semantic_nerf.py b/nerfactory/models/semantic_nerf.py index 8aa2a9923..8a84363ab 100644 --- a/nerfactory/models/semantic_nerf.py +++ b/nerfactory/models/semantic_nerf.py @@ -15,16 +15,19 @@ """ Semantic NeRF implementation. """ +from __future__ import annotations import torch from torch import nn from nerfactory.cameras.rays import RayBundle +from nerfactory.configs import base as cfg from nerfactory.dataloaders.structs import Semantics from nerfactory.fields.modules.encoding import NeRFEncoding from nerfactory.fields.modules.field_heads import FieldHeadNames from nerfactory.fields.nerf_field import NeRFField from nerfactory.fields.semantic_nerf_field import SemanticNerfField +from nerfactory.models.modules.scene_colliders import AABBBoxCollider from nerfactory.models.vanilla_nerf import NeRFModel from nerfactory.renderers.renderers import SemanticRenderer from nerfactory.utils import misc, writer @@ -33,10 +36,10 @@ from nerfactory.utils import misc, writer class SemanticNerfModel(NeRFModel): """Semantic-NeRF model""" - def __init__(self, semantics: Semantics, **kwargs) -> None: + def __init__(self, config: cfg.ModelConfig, semantics: Semantics, **kwargs) -> None: self.stuff_classes = semantics.stuff_classes self.stuff_colors = semantics.stuff_colors - super().__init__(**kwargs) + super().__init__(config=config, **kwargs) def populate_fields(self): """Set the fields.""" @@ -55,9 +58,17 @@ class SemanticNerfModel(NeRFModel): def populate_misc_modules(self): super().populate_misc_modules() + + # renderers self.renderer_semantic = SemanticRenderer() + + # losses self.cross_entropy_loss = nn.CrossEntropyLoss(reduction="mean") + # colliders + if self.config.enable_collider: + self.collider = AABBBoxCollider(scene_bounds=self.scene_bounds) + def get_outputs(self, ray_bundle: RayBundle): # uniform sampling ray_samples_uniform = self.sampler_uniform(ray_bundle) diff --git a/nerfactory/models/vanilla_nerf.py b/nerfactory/models/vanilla_nerf.py index 09eb3cc60..332630ec9 100644 --- a/nerfactory/models/vanilla_nerf.py +++ b/nerfactory/models/vanilla_nerf.py @@ -16,22 +16,24 @@ Implementation of vanilla nerf. """ +from __future__ import annotations -from typing import Dict, List, Optional +from typing import Dict, List import torch -from omegaconf import DictConfig from torch.nn import Parameter from torchmetrics import PeakSignalNoiseRatio from torchmetrics.functional import structural_similarity_index_measure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from nerfactory.cameras.rays import RayBundle +from nerfactory.configs import base as cfg from nerfactory.fields.modules.encoding import NeRFEncoding from nerfactory.fields.modules.field_heads import FieldHeadNames from nerfactory.fields.nerf_field import NeRFField from nerfactory.models.base import Model from nerfactory.models.modules.ray_sampler import PDFSampler, UniformSampler +from nerfactory.models.modules.scene_colliders import NearFarCollider from nerfactory.optimizers.loss import MSELoss from nerfactory.renderers.renderers import ( AccumulationRenderer, @@ -56,23 +58,14 @@ class NeRFModel(Model): def __init__( self, - near_plane: float = 2.0, - far_plane: float = 6.0, - num_coarse_samples: int = 64, - num_importance_samples: int = 128, - enable_density_field: bool = False, - density_field_config: Optional[DictConfig] = None, + config: cfg.ModelConfig, **kwargs, ) -> None: - self.near_plane = near_plane - self.far_plane = far_plane - self.num_coarse_samples = num_coarse_samples - self.num_importance_samples = num_importance_samples self.field_coarse = None self.field_fine = None + super().__init__( - enable_density_field=enable_density_field, - density_field_config=density_field_config, + config=config, **kwargs, ) @@ -106,8 +99,10 @@ class NeRFModel(Model): def populate_misc_modules(self): # samplers - self.sampler_uniform = UniformSampler(num_samples=self.num_coarse_samples, density_field=self.density_field) - self.sampler_pdf = PDFSampler(num_samples=self.num_importance_samples, density_field=self.density_field) + self.sampler_uniform = UniformSampler( + num_samples=self.config.num_coarse_samples, density_field=self.density_field + ) + self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples, density_field=self.density_field) # renderers self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) @@ -122,6 +117,12 @@ class NeRFModel(Model): self.ssim = structural_similarity_index_measure self.lpips = LearnedPerceptualImagePatchSimilarity() + # colliders + if self.config.enable_collider: + self.collider = NearFarCollider( + near_plane=self.config.collider_params["near_plane"], far_plane=self.config.collider_params["far_plane"] + ) + def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} if self.field_coarse is None or self.field_fine is None: @@ -191,14 +192,14 @@ class NeRFModel(Model): depth_coarse = visualization.apply_depth_colormap( outputs["depth_coarse"], accumulation=outputs["accumulation_coarse"], - near_plane=self.near_plane, - far_plane=self.far_plane, + near_plane=self.config.collider_params["near_plane"], + far_plane=self.config.collider_params["far_plane"], ) depth_fine = visualization.apply_depth_colormap( outputs["depth_fine"], accumulation=outputs["accumulation_fine"], - near_plane=self.near_plane, - far_plane=self.far_plane, + near_plane=self.config.collider_params["near_plane"], + far_plane=self.config.collider_params["far_plane"], ) combined_rgb = torch.cat([image, rgb_coarse, rgb_fine], dim=1) diff --git a/nerfactory/optimizers/optimizers.py b/nerfactory/optimizers/optimizers.py index b358ee2a7..ec8db3213 100644 --- a/nerfactory/optimizers/optimizers.py +++ b/nerfactory/optimizers/optimizers.py @@ -15,15 +15,13 @@ """ Optimizers class. """ +from __future__ import annotations from typing import Any, Dict, List, Union -import numpy as np -from hydra.utils import instantiate from omegaconf import DictConfig from torch.cuda.amp.grad_scaler import GradScaler from torch.nn.parameter import Parameter -from torch.optim.lr_scheduler import LambdaLR from nerfactory.utils import writer @@ -43,39 +41,6 @@ def setup_optimizers( return Optimizers(config, param_groups) -class ExponentialDecaySchedule(LambdaLR): - """Exponential learning rate decay function. - See https://github.com/google-research/google-research/blob/ - fd2cea8cdd86b3ed2c640cbe5561707639e682f3/jaxnerf/nerf/utils.py#L360 - for details. - - Args: - optimizer: The optimizer to update. - lr_init: The initial learning rate. - lr_final: The final learning rate. - max_steps: The maximum number of steps. - lr_delay_steps: The number of steps to delay the learning rate. Default is 0. - lr_delay_mult: The multiplier for the learning rate after the delay. Default is 1.0. - """ - - def __init__(self, optimizer, lr_init, lr_final, max_steps, lr_delay_steps=0, lr_delay_mult=1.0) -> None: - def func(step): - if lr_delay_steps > 0: - delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin( - 0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1) - ) - else: - delay_rate = 1.0 - t = np.clip(step / max_steps, 0, 1) - log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t) - multiplier = ( - log_lerp / lr_init - ) # divided by lr_init because the multiplier is with the initial learning rate - return delay_rate * multiplier - - super().__init__(optimizer, lr_lambda=func) - - class Optimizers: """A set of optimizers. @@ -89,11 +54,11 @@ class Optimizers: self.optimizers = {} self.schedulers = {} for param_group_name, params in param_groups.items(): - lr_init = config[param_group_name].optimizer.lr - self.optimizers[param_group_name] = instantiate(config[param_group_name].optimizer, params=params) - if config[param_group_name].scheduler: - self.schedulers[param_group_name] = instantiate( - config[param_group_name].scheduler, optimizer=self.optimizers[param_group_name], lr_init=lr_init + lr_init = config[param_group_name]["optimizer"].lr + self.optimizers[param_group_name] = config[param_group_name]["optimizer"].setup(params=params) + if config[param_group_name]["scheduler"]: + self.schedulers[param_group_name] = config[param_group_name]["scheduler"].setup( + optimizer=self.optimizers[param_group_name], lr_init=lr_init ) def optimizer_step(self, param_group_name: str) -> None: diff --git a/nerfactory/optimizers/schedulers.py b/nerfactory/optimizers/schedulers.py new file mode 100644 index 000000000..a493a5622 --- /dev/null +++ b/nerfactory/optimizers/schedulers.py @@ -0,0 +1,51 @@ +# 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. + +"""Scheduler Classes""" + +import numpy as np +from torch.optim.lr_scheduler import LambdaLR + + +class ExponentialDecaySchedule(LambdaLR): + """Exponential learning rate decay function. + See https://github.com/google-research/google-research/blob/ + fd2cea8cdd86b3ed2c640cbe5561707639e682f3/jaxnerf/nerf/utils.py#L360 + for details. + + Args: + optimizer: The optimizer to update. + lr_init: The initial learning rate. + lr_final: The final learning rate. + max_steps: The maximum number of steps. + lr_delay_steps: The number of steps to delay the learning rate. Default is 0. + lr_delay_mult: The multiplier for the learning rate after the delay. Default is 1.0. + """ + + def __init__(self, optimizer, lr_init, lr_final, max_steps, lr_delay_steps=0, lr_delay_mult=1.0) -> None: + def func(step): + if lr_delay_steps > 0: + delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin( + 0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1) + ) + else: + delay_rate = 1.0 + t = np.clip(step / max_steps, 0, 1) + log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t) + multiplier = ( + log_lerp / lr_init + ) # divided by lr_init because the multiplier is with the initial learning rate + return delay_rate * multiplier + + super().__init__(optimizer, lr_lambda=func) diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py index 67e2bff6a..8d44e4789 100644 --- a/nerfactory/pipelines/base.py +++ b/nerfactory/pipelines/base.py @@ -15,18 +15,18 @@ """ Abstracts for the Pipeline class. """ +from __future__ import annotations from abc import abstractmethod -from pydoc import locate from typing import Any, Dict, List, Optional from torch import nn from torch.nn import Parameter -from nerfactory.dataloaders.base import Dataloader, setup_dataloader -from nerfactory.models.base import Model, setup_model +from nerfactory.configs import base as cfg +from nerfactory.dataloaders.base import Dataloader +from nerfactory.models.base import Model from nerfactory.utils import profiler -from nerfactory.utils.config import PipelineConfig class Pipeline(nn.Module): @@ -70,10 +70,13 @@ class Pipeline(nn.Module): self.model (Model): The model that will be used """ - def __init__(self, dataloader: Dataloader, model: Model): + def __init__(self, config: cfg.PipelineConfig, device: str, test_mode: bool = False): super().__init__() - self.dataloader: Dataloader = dataloader - self.model: Model = model + self.dataloader: Dataloader = config.dataloader.setup(device=device, test_mode=test_mode) + self.dataloader.to(device) + # TODO(ethan): get rid of scene_bounds from the model + self.model: Model = config.model.setup(scene_bounds=self.dataloader.train_datasetinputs.scene_bounds) + self.model.to(device) @property def device(self): @@ -123,19 +126,3 @@ class Pipeline(nn.Module): model_params = self.model.get_param_groups() # TODO(ethan): assert that key names don't overlap return {**dataloader_params, **model_params} - - -@profiler.time_function -def setup_pipeline(config: PipelineConfig, device: str, test_mode=False) -> Pipeline: - """Setup the pipeline. The dataset inputs should be set with the training data. - - Args: - config: The pipeline config. - """ - # dataset_inputs - dataloader: Dataloader = setup_dataloader(config.dataloader, device=device, test_mode=test_mode) - # TODO(ethan): get rid of scene_bounds from the model - model: Model = setup_model(config.model, scene_bounds=dataloader.train_datasetinputs.scene_bounds, device=device) - pipeline_class = locate(config._target_) # pylint: disable=protected-access - pipeline = pipeline_class(dataloader=dataloader, model=model) # type: ignore - return pipeline diff --git a/nerfactory/utils/config.py b/nerfactory/utils/config.py deleted file mode 100644 index 62abd2d14..000000000 --- a/nerfactory/utils/config.py +++ /dev/null @@ -1,161 +0,0 @@ -# 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. - -"""Structured config classes""" - -from dataclasses import dataclass -from typing import Any, Dict, List, Optional - -from omegaconf import MISSING, DictConfig - - -@dataclass -class MachineConfig: - """Configuration of machine setup""" - - seed: int = MISSING - num_gpus: int = MISSING - num_machines: int = MISSING - machine_rank: int = MISSING - dist_url: str = MISSING - - -@dataclass -class LoggingConfig: - """Configuration of loggers and profilers""" - - steps_per_log: int = MISSING - max_buffer_size: int = MISSING - writer: Dict[str, Any] = MISSING - enable_profiler: bool = MISSING - - -@dataclass -class ResumeTrainConfig: - """Configuration for loading previous checkpoints""" - - load_dir: Optional[str] = None - load_step: Optional[int] = None - - -@dataclass -class TrainerConfig: - """Configuration for training regimen""" - - model_dir: str = MISSING - steps_per_save: int = MISSING - steps_per_test: int = MISSING - max_num_iterations: int = MISSING - mixed_precision: bool = MISSING - resume_train: ResumeTrainConfig = MISSING - - -@dataclass -class DataloaderConfig: - """Configuration for train/eval datasets""" - - _target_: str = MISSING - image_dataset_type: Optional[str] = "rgb" - train_dataset: Dict[str, Any] = MISSING - train_num_rays_per_batch: int = MISSING - train_num_images_to_sample_from: int = MISSING - eval_dataset: Optional[Dict[str, Any]] = None - eval_image_indices: Optional[List[int]] = None - eval_num_rays_per_chunk: int = MISSING - - -@dataclass -class ModelConfig: - """Configuration for graph instantiation""" - - _target_: str = MISSING - enable_collider: Optional[bool] = True - collider_config: Dict[str, Any] = MISSING - num_coarse_samples: int = MISSING - num_importance_samples: int = MISSING - loss_coefficients: Dict[str, Any] = MISSING - # additional optional parameters here - field_implementation: Optional[str] = "torch" - enable_density_field: Optional[bool] = False - density_field_config: Dict[str, Any] = MISSING - - -@dataclass -class PipelineConfig: - """Configuration for pipeline instantiation""" - - _target_: str = MISSING - dataloader: DataloaderConfig = MISSING - model: ModelConfig = MISSING - - -@dataclass -class ViewerConfig: - """Configuration for viewer instantiation""" - - enable: bool = MISSING - zmq_url: str = MISSING - min_render_image_height: int = MISSING - max_render_image_height: int = MISSING - num_rays_per_chunk: int = MISSING - - -@dataclass -class Config: - """Full config contents""" - - machine: MachineConfig = MISSING - logging: LoggingConfig = MISSING - trainer: TrainerConfig = MISSING - experiment_name: str = MISSING - method_name: str = MISSING - optimizers: Dict[str, Any] = MISSING - viewer: ViewerConfig = MISSING - pipeline: PipelineConfig = MISSING - # additional optional parameters here - hydra: Optional[Dict[str, Any]] = None - - -def setup_config(config: DictConfig) -> Config: - """helper that creates a typed config from the DictConfig - - Args: - config (DictConfig): configuration to convert - - Returns: - Config: typed version of the input configuration - """ - machine = MachineConfig(**config.machine) - logging = LoggingConfig(**config.logging) - trainer = TrainerConfig(**config.trainer) - experiment_name = config.experiment_name - method_name = config.method_name - dataloader = DataloaderConfig(**config.pipeline.dataloader) - model = ModelConfig(**config.pipeline.model) - pipeline = PipelineConfig( - _target_=config.pipeline._target_, dataloader=dataloader, model=model # pylint:disable=protected-access - ) - optimizers = config.optimizers - viewer = ViewerConfig(**config.viewer) - - return Config( - machine=machine, - logging=logging, - trainer=trainer, - experiment_name=experiment_name, - method_name=method_name, - pipeline=pipeline, - optimizers=optimizers, - viewer=viewer, - ) diff --git a/nerfactory/utils/misc.py b/nerfactory/utils/misc.py index 3ff9ca2ac..5f4a5d0e9 100644 --- a/nerfactory/utils/misc.py +++ b/nerfactory/utils/misc.py @@ -19,23 +19,9 @@ Miscellaneous helper code. import hashlib import json from math import floor, log -from pydoc import locate from typing import Any, Callable, Dict, Union import torch -from omegaconf import DictConfig - - -class DotDict(dict): - """ - dot.notation access to dictionary attributes - """ - - def __getattr__(self, attr): - return self[attr] - - __setattr__ = dict.__setitem__ # type: ignore - __delattr__ = dict.__delitem__ # type: ignore def get_dict_to_torch(stuff, device: Union[torch.device, str] = "cpu"): @@ -74,23 +60,6 @@ def get_masked_dict(d, mask): return masked_dict -def instantiate_from_dict_config(dict_config: Any, **kwargs): - """Our version of hydra's instantiate function. - - Args: - dict_config: DictConfig object to instantiate. It can be a dataclass or a dict but must have a `_target_` field. - """ - dict_config_kwargs = {k: v for k, v in dict_config.items() if k != "_target_"} - uninstantiated_class = locate(dict_config._target_) # pylint: disable=protected-access - assert ( - uninstantiated_class is not None - ), f"Could not find class {dict_config._target_}" # pylint: disable=protected-access - all_kwargs = dict_config_kwargs - all_kwargs.update(kwargs) - instantiated_class = uninstantiated_class(**all_kwargs) # type: ignore - return instantiated_class - - def get_hash_str_from_dict(dictionary: Dict[str, Any]) -> str: """MD5 hash of a dictionary. Based on https://www.doc.ic.ac.uk/~nuric/coding/how-to-hash-a-dictionary-in-python.html""" @@ -150,7 +119,7 @@ def human_format(num): return f"{(num / k**magnitude):.2f} {units[magnitude]}" -def scale_dict(dictionary: Dict[Any, Any], coefficients: DictConfig) -> Dict[Any, Any]: +def scale_dict(dictionary: Dict[Any, Any], coefficients: Dict[str, float]) -> Dict[Any, Any]: """Scale a dictionary in-place given a coefficients dictionary. Args: diff --git a/nerfactory/utils/profiler.py b/nerfactory/utils/profiler.py index 41728ed2b..ee69c2bf5 100644 --- a/nerfactory/utils/profiler.py +++ b/nerfactory/utils/profiler.py @@ -15,12 +15,14 @@ """ Profiler base class and functionality """ +from __future__ import annotations + import logging import time from typing import Callable +from nerfactory.configs import base as cfg from nerfactory.utils import comms -from nerfactory.utils.config import LoggingConfig from nerfactory.utils.decorators import ( check_main_thread, check_profiler_enabled, @@ -44,13 +46,13 @@ def time_function(func: Callable) -> Callable: return wrapper -def flush_profiler(config: LoggingConfig): +def flush_profiler(config: cfg.LoggingConfig): """Method that checks if profiler is enabled before flushing""" if config.enable_profiler and PROFILER: PROFILER[0].print_profile() -def setup_profiler(config: LoggingConfig): +def setup_profiler(config: cfg.LoggingConfig): """Initialization of profilers""" if comms.is_main_process(): PROFILER.append(Profiler(config)) @@ -60,7 +62,7 @@ def setup_profiler(config: LoggingConfig): class Profiler: """Profiler class""" - def __init__(self, config: LoggingConfig): + def __init__(self, config: cfg.LoggingConfig): self.config = config self.profiler_dict = {} diff --git a/nerfactory/utils/writer.py b/nerfactory/utils/writer.py index ae9695b80..64ee430f8 100644 --- a/nerfactory/utils/writer.py +++ b/nerfactory/utils/writer.py @@ -15,24 +15,23 @@ """ Generic Writer class """ - +from __future__ import annotations import enum -import os -import sys +import logging from abc import abstractmethod +from pathlib import Path from time import time -from typing import Dict +from typing import Dict, Optional import imageio import numpy as np import torch import wandb -from omegaconf import ListConfig from torch.utils.tensorboard import SummaryWriter from torchtyping import TensorType -from nerfactory.utils.config import LoggingConfig +from nerfactory.configs import base as cfg from nerfactory.utils.decorators import check_main_thread, decorate_all from nerfactory.utils.misc import human_format @@ -138,16 +137,13 @@ def write_out_storage(): @check_main_thread -def setup_event_writers(config: LoggingConfig, max_iter: int) -> None: +def setup_event_writers(config: cfg.LoggingConfig, max_iter: int) -> None: """Initialization of all event writers specified in config""" - for writer_type in config.writer: - writer_class = getattr(sys.modules[__name__], writer_type) - writer_config = config.writer[writer_type] - if writer_type == "LocalWriter": - curr_writer = writer_class(writer_config.log_dir, writer_config.stats_to_track, writer_config.max_log_size) - else: - curr_writer = writer_class(writer_config.log_dir) - EVENT_WRITERS.append(curr_writer) + for writer_type_config in config.writer: + if writer_type_config.enable: + curr_writer = writer_type_config.setup() + EVENT_WRITERS.append(curr_writer) + logging.info("logging info to: %s", writer_type_config.log_dir) ## configure all the global buffer basic information GLOBAL_BUFFER["max_iter"] = max_iter @@ -159,7 +155,7 @@ def setup_event_writers(config: LoggingConfig, max_iter: int) -> None: class Writer: """Writer class""" - def __init__(self, log_dir: str): + def __init__(self, log_dir: Optional[Path]): self.log_dir = log_dir @abstractmethod @@ -228,9 +224,9 @@ class TimeWriter: class WandbWriter(Writer): """WandDB Writer Class""" - def __init__(self, log_dir: str): - super().__init__(log_dir) - wandb.init(dir=log_dir) + def __init__(self, config: cfg.WandbWriterConfig): + super().__init__(config.log_dir) + wandb.init(dir=config.log_dir) def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: """_summary_ @@ -256,9 +252,9 @@ class WandbWriter(Writer): class TensorboardWriter(Writer): """Tensorboard Writer Class""" - def __init__(self, log_dir: str): - super().__init__(log_dir) - self.tb_writer = SummaryWriter(log_dir=self.log_dir) + def __init__(self, config: cfg.TensorboardWriterConfig): + super().__init__(config.log_dir) + self.tb_writer = SummaryWriter(log_dir=config.log_dir) def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: """_summary_ @@ -313,23 +309,23 @@ def _format_time(seconds): class LocalWriter(Writer): """Local Writer Class""" - def __init__(self, log_dir: str, stats_to_track: ListConfig, max_log_size: int = 0): + def __init__(self, config: cfg.LocalWriterConfig): """ Args: stats_to_track (ListConfig): the names of stats that should be logged. max_log size (int): max number of lines that will be logged to teminal. """ - super().__init__(log_dir) - self.stats_to_track = [EventName[name].value for name in stats_to_track] - self.max_log_size = max_log_size + super().__init__(config.log_dir) + self.stats_to_track = [name.value for name in config.stats_to_track] + self.max_log_size = config.max_log_size self.keys = set() self.past_mssgs = ["", ""] self.has_printed = False def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: - if name in self.stats_to_track: + if name in self.stats_to_track and self.log_dir: image = to8b(image) - image_path = os.path.join(self.log_dir, f"{name}.jpg") + image_path = self.log_dir / f"{name}.jpg" imageio.imwrite(image_path, np.uint8(image.cpu().numpy() * 255.0)) def write_scalar(self, name: str, scalar: float, step: int) -> None: diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index 62b03415d..1485ea2cc 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/nerfactory/viewer/server/viewer_utils.py @@ -27,11 +27,11 @@ import torch from nerfactory.cameras.cameras import get_camera, get_intrinsics_from_intrinsics_matrix from nerfactory.cameras.rays import RayBundle +from nerfactory.configs import base as cfg from nerfactory.dataloaders.image_dataset import ImageDataset from nerfactory.dataloaders.structs import DatasetInputs from nerfactory.models.base import Model from nerfactory.utils import profiler, visualization -from nerfactory.utils.config import ViewerConfig from nerfactory.utils.decorators import check_visualizer_enabled, decorate_all from nerfactory.utils.writer import GLOBAL_BUFFER, EventName from nerfactory.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_h @@ -184,7 +184,7 @@ class VisualizerState: config: viewer setup configuration """ - def __init__(self, config: ViewerConfig): + def __init__(self, config: cfg.ViewerConfig): self.config = config self.vis = None diff --git a/scripts/benchmarking/launch_eval_blender.sh b/scripts/benchmarking/launch_eval_blender.sh index 53d6c836e..ccf153ca1 100755 --- a/scripts/benchmarking/launch_eval_blender.sh +++ b/scripts/benchmarking/launch_eval_blender.sh @@ -59,11 +59,9 @@ for dataset in ${DATASETS[@]}; do export CUDA_VISIBLE_DEVICES=${GPU_IDX[$idx]} base_config_name=${config_name/"graph_"/""} checkpoint_dir="${output_dir}/blender_${dataset}_${month}-${date}-${year}/${base_config_name}/${year}-${month}-${date}_${seconds}/nerfactory_models/" - python scripts/run_eval.py --config-name ${config_name} \ - --checkpoint-dir ${checkpoint_dir} \ - --rendered-output-name rgb \ - --output-filename ${output_dir}/${base_config_name}/blender_${dataset}_${month}-${date}-${year}_${seconds}.json \ - --config-overrides pipeline.dataloader.train_dataset.data_directory=data/blender/${dataset} & + python scripts/run_eval.py ${config_name} \ + --eval.checkpoint-dir=${checkpoint_dir} \ + --eval.output-filename=${output_dir}/${base_config_name}/blender_${dataset}_${month}-${date}-${year}_${seconds}.json & echo "Launched ${checkpoint_dir} on gpu ${GPU_IDX[$idx]}" # update gpu diff --git a/scripts/benchmarking/launch_train_blender.sh b/scripts/benchmarking/launch_train_blender.sh index 1bbe9a50c..80c91018d 100755 --- a/scripts/benchmarking/launch_train_blender.sh +++ b/scripts/benchmarking/launch_train_blender.sh @@ -3,7 +3,7 @@ helpFunction_launch_train() { echo "Usage: $0 -c config_name -g gpu_list" - echo -e "\t-c name of config to benchmark (e.g. graph_mipnerf, graph_instant_ngp)" + echo -e "\t-c name of config to benchmark (e.g. mipnerf, instant_ngp)" echo -e "\t-g list of space-separated gpu numbers to launch train on (e.g. 0 2 4 5)" exit 1 # Exit program after printing help } @@ -53,15 +53,14 @@ len=${#GPU_IDX[@]} for dataset in ${DATASETS[@]}; do export CUDA_VISIBLE_DEVICES=${GPU_IDX[$idx]} python scripts/run_train.py \ - --config-name ${config_name} \ - '~logging.writer.LocalWriter' \ - pipeline.dataloader.train_dataset.data_directory=data/blender/${dataset} \ - experiment_name=blender_${dataset}_${tag} \ - trainer.model_dir=nerfactory_models/ \ - trainer.steps_per_save=25000 \ - trainer.max_num_iterations=2000000 \ - viewer.enable=False \ - logging.enable_profiler=False & + ${config_name} \ + --logging.writer.2.no-enable \ + --pipeline.dataloader.train-dataset.data-directory=data/blender/${dataset} \ + --experiment-name=blender_${dataset}_${tag} \ + --trainer.model-dir=nerfactory_models/ \ + --trainer.steps-per-save=16000 \ + --trainer.max-num-iterations=16500 \ + --logging.no-enable-profiler & echo "Launched ${config_name} ${dataset} on gpu ${GPU_IDX[$idx]}, ${tag}" # update gpu diff --git a/scripts/run_eval.py b/scripts/run_eval.py index 689a081c3..9935af834 100644 --- a/scripts/run_eval.py +++ b/scripts/run_eval.py @@ -1,16 +1,15 @@ """ run_eval.py """ -import enum import json +import logging import os -from typing import Dict, List, Optional +import pickle +from pathlib import Path +from typing import Dict -import dcargs import mediapy as media import torch -from hydra import compose, initialize -from omegaconf import DictConfig from tqdm import tqdm from nerfactory.cameras.camera_paths import ( @@ -18,12 +17,14 @@ from nerfactory.cameras.camera_paths import ( get_interpolated_camera_path, get_spiral_path, ) -from nerfactory.pipelines.base import Pipeline, setup_pipeline -from nerfactory.utils import io as io_utils -from nerfactory.utils.config import setup_config +from nerfactory.configs import base as cfg +from nerfactory.configs.utils import cli_from_base_configs +from nerfactory.pipelines.base import Pipeline from nerfactory.utils.misc import human_format from nerfactory.utils.writer import TimeWriter +logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.INFO) + def _update_avg(prev_avg: float, new_val: float, step: int) -> float: """helper to calculate the running average @@ -39,7 +40,7 @@ def _update_avg(prev_avg: float, new_val: float, step: int) -> float: return (step * prev_avg + new_val) / (step + 1) -def _load_checkpoint(config: DictConfig, pipeline: Pipeline) -> str: +def _load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path: """Helper function to load checkpointed pipeline Args: @@ -53,8 +54,8 @@ def _load_checkpoint(config: DictConfig, pipeline: Pipeline) -> str: load_step = sorted(int(x[x.find("-") + 1 : x.find(".")]) for x in os.listdir(config.load_dir))[-1] else: load_step = config.load_step - load_path = os.path.join(config.load_dir, f"step-{load_step:09d}.ckpt") - assert os.path.exists(load_path), f"Checkpoint {load_path} does not exist" + load_path = config.load_dir / f"step-{load_step:09d}.ckpt" + assert load_path.exists(), f"Checkpoint {load_path} does not exist" loaded_state = torch.load(load_path, map_location="cpu") pipeline.load_pipeline(loaded_state["pipeline"]) print(f"done loading checkpoint from {load_path}") @@ -88,7 +89,7 @@ def render_stats_dict(pipeline: Pipeline) -> Dict[str, float]: def render_trajectory_video( pipeline: Pipeline, camera_path: CameraPath, - output_filename: str, + output_filename: Path, rendered_output_name: str, rendered_resolution_scaling_factor: float = 1.0, num_rays_per_chunk: int = 4096, @@ -120,30 +121,7 @@ def render_trajectory_video( media.write_video(output_filename, images, fps=fps) -class MethodType(enum.Enum): - """Enum for the method type.""" - - PSNR = enum.auto() - TRAJ = enum.auto() - - -class TrajectoryType(enum.Enum): - """Enum for the trajectory type.""" - - SPIRAL = enum.auto() - INTERP = enum.auto() - - -def main( - config_name: str, - checkpoint_dir: str, - rendered_output_name: str, - method: MethodType = MethodType.PSNR, - traj: TrajectoryType = TrajectoryType.SPIRAL, - output_filename: str = "output.json", - rendered_resolution_scaling_factor: float = 1.0, - config_overrides: Optional[List[str]] = None, -): +def main(config: cfg.Config): """Evaluate trained model. This evaluation can either render a trajectory or compute the eval psnr. Args: @@ -157,25 +135,21 @@ def main( Defaults to 1.0. config_overrides: List of strings to override config values. """ + assert config.eval is not None, "Eval config needs to be set" + assert config.eval.checkpoint_dir is not None, "Missing eval checkpoint dir: --eval.checkpoint-dir" - config_path = "../configs" - initialize(version_base="1.2", config_path=config_path) - config_overrides = config_overrides or [] - config = compose(config_name, overrides=config_overrides) - - config.trainer.resume_train.load_dir = checkpoint_dir + config.trainer.load_dir = config.eval.checkpoint_dir config.pipeline.dataloader.eval_image_indices = None - config = setup_config(config) # converting to typed config device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # setup pipeline (which includes the dataloaders) - pipeline = setup_pipeline(config.pipeline, device=device, test_mode=True) + pipeline = config.pipeline.setup(device=device, test_mode=True) pipeline.eval() # load checkpointed information - checkpoint_name = _load_checkpoint(config.trainer.resume_train, pipeline) + checkpoint_name = _load_checkpoint(config.trainer, pipeline) - if method == MethodType.PSNR: - assert output_filename.endswith(".json") + if config.eval.method == cfg.MethodType.PSNR: + assert config.eval.output_filename.suffix == ".json" stats_dict = render_stats_dict(pipeline) avg_psnr = stats_dict["avg psnr"] avg_rays_per_sec = stats_dict["avg rays per sec"] @@ -184,36 +158,53 @@ def main( print(f"Avg. Rays / Sec: {human_format(avg_rays_per_sec)}") print(f"Avg. FPS: {avg_fps:0.4f}") # save output to some file - io_utils.make_dir(output_filename) + config.eval.output_filename.parent.mkdir(parents=True, exist_ok=True) benchmark_info = { - "config_name": config_name, - "checkpoint": checkpoint_name, + "experiment_name": config.experiment_name, + "method_name": config.method_name, + "checkpoint": str(checkpoint_name), "results": stats_dict, } - with open(output_filename, "w", encoding="utf8") as f: - json.dump(benchmark_info, f, indent=2) - print(f"Saved results to: {output_filename}") + with open(config.eval.output_filename, "w", encoding="utf8") as output_f: + json.dump(benchmark_info, output_f, indent=2) + print(f"Saved results to: {config.eval.output_filename}") + + elif config.eval.method == cfg.MethodType.TRAJ: + assert config.eval.output_filename.suffix == ".mp4" + assert config.eval.rendered_output_name is not None, "Missing eval output type: --eval.rendered-output-name" - elif method == MethodType.TRAJ: - assert output_filename.endswith(".mp4") # TODO(ethan): use camera information from parsing args - if traj == TrajectoryType.SPIRAL: + if config.eval.traj == cfg.TrajectoryType.SPIRAL: camera_start = pipeline.dataloader.eval_dataloader.get_camera(image_idx=0) # TODO(ethan): pass in the up direction of the camera camera_path = get_spiral_path(camera_start, steps=30, radius=0.1) - elif traj == TrajectoryType.INTERP: + elif config.eval.traj == cfg.TrajectoryType.INTERP: camera_start = pipeline.dataloader.eval_dataloader.get_camera(image_idx=0) camera_end = pipeline.dataloader.eval_dataloader.get_camera(image_idx=10) camera_path = get_interpolated_camera_path(camera_start, camera_end, steps=30) + else: + raise NotImplementedError + render_trajectory_video( pipeline, camera_path, - output_filename=output_filename, - rendered_output_name=rendered_output_name, - rendered_resolution_scaling_factor=rendered_resolution_scaling_factor, + output_filename=config.eval.output_filename, + rendered_output_name=config.eval.rendered_output_name, + rendered_resolution_scaling_factor=config.eval.rendered_resolution_scaling_factor, num_rays_per_chunk=pipeline.dataloader.eval_num_rays_per_chunk, ) if __name__ == "__main__": - dcargs.cli(main) + from nerfactory.configs.base_configs import base_configs + + instantiated_config = cli_from_base_configs(base_configs, eval_mode=True) + if instantiated_config.eval.load_config: + logging.info(f"Loading pre-set config to: {instantiated_config.eval.load_config}") + with open(instantiated_config.eval.load_config, "rb") as config_f: + train_instantiated_config = pickle.load(config_f) + train_instantiated_config.eval = instantiated_config.eval + instantiated_config = train_instantiated_config + logging.info("Printing current config setup") + print(instantiated_config) + main(instantiated_config) diff --git a/scripts/run_train.py b/scripts/run_train.py index edfde0ced..f7992fd8e 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -3,25 +3,22 @@ run_train_nerf.py """ import logging -import os +import pickle import random import socket import traceback from datetime import timedelta -from typing import Any, Callable +from typing import Any, Callable, Optional -import hydra import numpy as np import torch import torch.distributed as dist import torch.multiprocessing as mp -import yaml -from hydra.core.config_store import ConfigStore -from omegaconf import DictConfig +from nerfactory.configs import base as cfg +from nerfactory.configs.utils import cli_from_base_configs from nerfactory.engine.trainer import train_loop from nerfactory.utils import comms, profiler -from nerfactory.utils.config import Config, setup_config logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG) @@ -54,7 +51,7 @@ def _distributed_worker( num_gpus_per_machine: int, machine_rank: int, dist_url: str, - config: Config, + config: cfg.Config, timeout: timedelta = DEFAULT_TIMEOUT, ) -> Any: """Spawned distributed worker that handles the initialization of process group and handles the @@ -89,8 +86,7 @@ def _distributed_worker( timeout=timeout, ) except Exception as e: - logger = logging.getLogger(__name__) - logger.error("Process group URL: %s", dist_url) + logging.error("Process group URL: %s", dist_url) raise e assert comms._LOCAL_PROCESS_GROUP is None # pylint: disable=protected-access @@ -118,7 +114,7 @@ def launch( num_machines: int = 1, machine_rank: int = 0, dist_url: str = "auto", - config: Config = None, + config: Optional[cfg.Config] = None, timeout: timedelta = DEFAULT_TIMEOUT, ) -> None: """Function that spawns muliple processes to call on main_func @@ -153,8 +149,7 @@ def launch( port = _find_free_port() dist_url = f"tcp://127.0.0.1:{port}" if num_machines > 1 and dist_url.startswith("file://"): - logger = logging.getLogger(__name__) - logger.warning("file:// is not a reliable init_method in multi-machine jobs. Prefer tcp://") + logging.warning("file:// is not a reliable init_method in multi-machine jobs. Prefer tcp://") process_context = mp.spawn( _distributed_worker, @@ -173,34 +168,18 @@ def launch( try: process_context.join() except KeyboardInterrupt: - logger = logging.getLogger(__name__) for i, process in enumerate(process_context.processes): if process.is_alive(): - logger.info("Terminating process %s", str(i)) + logging.info("Terminating process %s", str(i)) process.terminate() process.join() - logger.info("Process %s finished", str(i)) + logging.info("Process %s finished", str(i)) finally: profiler.flush_profiler(config.logging) -cs = ConfigStore.instance() -cs.store(name="graph_default", node=Config) - - -@hydra.main(version_base="1.2", config_path="../configs", config_name="graph_default.yaml") -def main(config: DictConfig): +def main(config: cfg.Config) -> None: """Main function.""" - config = setup_config(config) # converting to typed config - - unrolled_path = os.path.join(os.getcwd(), ".hydra/config.yaml") - if os.path.exists(unrolled_path): - with open(unrolled_path, encoding="utf8") as f: - unrolled_config = yaml.safe_load(f) - logger = logging.getLogger(__name__) - logger.info("Printing current config setup") - print(yaml.dump(unrolled_config, sort_keys=False, default_flow_style=False)) - launch( train_loop, config.machine.num_gpus, @@ -212,4 +191,23 @@ def main(config: DictConfig): if __name__ == "__main__": - main() # pylint: disable=no-value-for-parameter + from nerfactory.configs.base_configs import base_configs + + instantiated_config = cli_from_base_configs(base_configs) + if instantiated_config.trainer.load_config: + logging.info(f"Loading pre-set config to: {instantiated_config.trainer.load_config}") + with open(instantiated_config.trainer.load_config, "rb") as f: + instantiated_config = pickle.load(f) + # log and save config + logging.info("Printing current config setup") + print(instantiated_config) + config_pickle_path = instantiated_config.base_dir / "config.pkl" + config_json_path = instantiated_config.base_dir / "config.txt" + logging.info(f"Saving instance of config to: {config_pickle_path}") + logging.info(f"Saving human-readable config to: {config_json_path}") + config_pickle_path.parent.mkdir(parents=True, exist_ok=True) + with open(config_pickle_path, "wb") as f: + pickle.dump(instantiated_config, f) + with open(config_json_path, "w", encoding="utf8") as f: + f.write(str(instantiated_config)) + main(config=instantiated_config) diff --git a/tests/graphs/modules/test_ray_sampler.py b/tests/graphs/modules/test_ray_sampler.py index f349b9b75..e5d4bfd6b 100644 --- a/tests/graphs/modules/test_ray_sampler.py +++ b/tests/graphs/modules/test_ray_sampler.py @@ -1,6 +1,7 @@ """ Test samplers """ + import torch from nerfactory.cameras.rays import RayBundle @@ -16,8 +17,6 @@ from nerfactory.models.modules.scene_colliders import NearFarCollider def test_uniform_sampler(): """Test uniform sampler""" - near_plane = 2 - far_plane = 4 num_samples = 15 sampler = UniformSampler(num_samples=num_samples) @@ -25,7 +24,7 @@ def test_uniform_sampler(): directions = torch.ones_like(origins) radius = torch.ones((10, 1)) ray_bundle = RayBundle(origins=origins, directions=directions, pixel_area=radius) - collider = NearFarCollider(near_plane, far_plane) + collider = NearFarCollider(near_plane=2, far_plane=4) ray_bundle = collider(ray_bundle) ray_samples = sampler(ray_bundle) @@ -37,8 +36,6 @@ def test_uniform_sampler(): def test_lin_disp_sampler(): """Test linear in disparity sampler""" - near_plane = 2 - far_plane = 4 num_samples = 15 sampler = LinearDisparitySampler(num_samples=num_samples) @@ -46,7 +43,7 @@ def test_lin_disp_sampler(): directions = torch.ones_like(origins) radius = torch.ones((10, 1)) ray_bundle = RayBundle(origins=origins, directions=directions, pixel_area=radius) - collider = NearFarCollider(near_plane, far_plane) + collider = NearFarCollider(near_plane=2, far_plane=4) ray_bundle = collider(ray_bundle) ray_samples = sampler(ray_bundle) @@ -56,8 +53,6 @@ def test_lin_disp_sampler(): def test_sqrt_sampler(): """Test square root sampler""" - near_plane = 2 - far_plane = 4 num_samples = 15 sampler = SqrtSampler(num_samples=num_samples) @@ -65,7 +60,7 @@ def test_sqrt_sampler(): directions = torch.ones_like(origins) radius = torch.ones((10, 1)) ray_bundle = RayBundle(origins=origins, directions=directions, pixel_area=radius) - collider = NearFarCollider(near_plane, far_plane) + collider = NearFarCollider(near_plane=2, far_plane=4) ray_bundle = collider(ray_bundle) ray_samples = sampler(ray_bundle) @@ -75,8 +70,6 @@ def test_sqrt_sampler(): def test_log_sampler(): """Test log sampler""" - near_plane = 2 - far_plane = 4 num_samples = 15 sampler = LogSampler(num_samples=num_samples) @@ -84,7 +77,7 @@ def test_log_sampler(): directions = torch.ones_like(origins) radius = torch.ones((10, 1)) ray_bundle = RayBundle(origins=origins, directions=directions, pixel_area=radius) - collider = NearFarCollider(near_plane, far_plane) + collider = NearFarCollider(near_plane=2, far_plane=4) ray_bundle = collider(ray_bundle) ray_samples = sampler(ray_bundle) @@ -94,15 +87,13 @@ def test_log_sampler(): def test_pdf_sampler(): """Test pdf sampler""" - near_plane = 2 - far_plane = 4 num_samples = 15 origins = torch.zeros((10, 3)) directions = torch.ones_like(origins) radius = torch.ones((10, 1)) ray_bundle = RayBundle(origins=origins, directions=directions, pixel_area=radius) - collider = NearFarCollider(near_plane, far_plane) + collider = NearFarCollider(near_plane=2, far_plane=4) ray_bundle = collider(ray_bundle) uniform_sampler = UniformSampler(num_samples=num_samples) diff --git a/tests/test_run_train.py b/tests/test_run_train.py index ffcbef041..b8a8a7836 100644 --- a/tests/test_run_train.py +++ b/tests/test_run_train.py @@ -2,20 +2,21 @@ """ Default test to make sure train runs """ -import os + +from __future__ import annotations + +from pathlib import Path import pytest -from hydra import compose, initialize -from hydra.core.global_hydra import GlobalHydra -from omegaconf import DictConfig, open_dict +import nerfactory.configs.base as cfg +from nerfactory.configs.base_configs import base_configs from nerfactory.engine.trainer import train_loop -from nerfactory.utils.config import setup_config -BLACKLIST = ["graph_semantic_nerf.yaml", "graph_mipnerf_360.yaml", "graph_instant_ngp.yaml"] +BLACKLIST = ["base", "semantic_nerf", "mipnerf_360", "instant_ngp"] -def set_reduced_config(config: DictConfig): +def set_reduced_config(config: cfg.Config): """Reducing the config settings to speedup test""" config.machine.num_gpus = 0 config.trainer.max_num_iterations = 2 @@ -23,25 +24,21 @@ def set_reduced_config(config: DictConfig): # switch to using the vanilla ImageDataset class config.pipeline.dataloader.image_dataset_type = "rgb" - config.pipeline.dataloader.train_dataset = { - "_target_": "nerfactory.dataloaders.datasets.Blender", - "data_directory": "tests/data/lego_test", - "downscale_factor": 16, - } + config.pipeline.dataloader.train_dataset = cfg.BlenderDatasetConfig( + data_directory="tests/data/lego_test", downscale_factor=16 + ) config.pipeline.dataloader.train_num_images_to_sample_from = 1 config.pipeline.dataloader.train_num_rays_per_batch = 4 - with open_dict(config): - config.pipeline.dataloader.eval_dataset = { - "_target_": "nerfactory.dataloaders.datasets.Blender", - "data_directory": "tests/data/lego_test", - "downscale_factor": 16, - } + config.pipeline.dataloader.eval_dataset = cfg.BlenderDatasetConfig( + data_directory="tests/data/lego_test", downscale_factor=16 + ) # reduce model factors config.pipeline.model.num_coarse_samples = 4 config.pipeline.model.num_importance_samples = 4 # set logging to tmp - config.logging.writer.TensorboardWriter.log_dir = "/tmp/" + for writer in config.logging.writer: + writer.log_dir = Path("/tmp/") # remove viewer config.viewer.enable = False @@ -55,20 +52,17 @@ def set_reduced_config(config: DictConfig): @pytest.mark.filterwarnings("ignore::DeprecationWarning") def test_run_train(): """test run train script works properly""" - all_configs = [f for f in os.listdir("configs/") if f.endswith(".yaml")] # relative to run path - for config_path in all_configs: - print(f"testing run for: {config_path}") - if config_path in BLACKLIST: + all_config_names = base_configs.keys() + for config_name in all_config_names: + if config_name in BLACKLIST: + print("skipping", config_name) continue - initialize(version_base="1.2", config_path="../configs/") # relative to test path - config = compose(config_path) + print(f"testing run for: {config_name}") + config = base_configs[config_name] config = set_reduced_config(config) - config = setup_config(config) train_loop(local_rank=0, world_size=0, config=config) - GlobalHydra.instance().clear() - if __name__ == "__main__": test_run_train()