From fe5f46c0aa6825bdaaadfc248e855a85ce5cce74 Mon Sep 17 00:00:00 2001 From: Brent Yi Date: Mon, 29 Aug 2022 17:07:33 -0700 Subject: [PATCH] Config suggestions (#300) * Config suggestions * Add some TODO notes --- README.md | 2 +- docs/tutorials/quickstart.md | 2 +- environment/requirements.txt | 4 +- nerfactory/configs/base.py | 29 ---- nerfactory/configs/base_configs.py | 4 +- nerfactory/configs/utils.py | 11 +- scripts/benchmarking/launch_eval_blender.sh | 10 +- scripts/run_eval.py | 140 ++++++++++++-------- scripts/run_train.py | 31 +++-- 9 files changed, 115 insertions(+), 118 deletions(-) diff --git a/README.md b/README.md index 33f3c0ee1..af3b2a0e0 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 --checkpoint-dir=outputs/blender_lego/instant_ngp/2022-07-07_230905/checkpoints +python scripts/run_eval.py render-trajectory --load-config=outputs/blender_lego/instant_ngp/2022-07-07_230905/config.yml--traj=spiral --output-path=output.mp4 ``` #### 6. In-depth guide diff --git a/docs/tutorials/quickstart.md b/docs/tutorials/quickstart.md index 16137ca9a..49c53377a 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 --checkpoint-dir=outputs/blender_lego/instant_ngp/2022-07-07_230905/checkpoints +python scripts/run_eval.py render-trajectory --load-config=outputs/blender_lego/instant_ngp/2022-07-07_230905/config.yml --traj=spiral --output-path=output.mp4 ``` 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 3cef67ffb..64f2481c3 100644 --- a/environment/requirements.txt +++ b/environment/requirements.txt @@ -2,7 +2,7 @@ aiohttp aiortc av black[jupyter]==22.3.0 -dcargs>=0.2.5 +dcargs>=0.2.6 ninja functorch==0.1.1 h5py>=2.9.0 @@ -30,4 +30,4 @@ torchtyping torchvision==0.12.0 tqdm u-msgpack-python>=2.4.1 -wandb \ No newline at end of file +wandb diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index d4b84eef4..f40df99b8 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -14,7 +14,6 @@ """Base Configs""" -import enum from dataclasses import dataclass from datetime import datetime from pathlib import Path @@ -320,33 +319,6 @@ class SchedulerConfig(InstantiateConfig): 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""" @@ -367,7 +339,6 @@ class Config(PrintableConfig): } ) viewer: ViewerConfig = ViewerConfig() - eval: Optional[EvalConfig] = None def __post_init__(self): """Convert logging directories to more specific filepaths""" diff --git a/nerfactory/configs/base_configs.py b/nerfactory/configs/base_configs.py index e5bd22e89..45ffc7c0c 100644 --- a/nerfactory/configs/base_configs.py +++ b/nerfactory/configs/base_configs.py @@ -16,6 +16,8 @@ Put all the method implementations in one location. """ +from typing import Dict + from nerfactory.configs.base import ( BlenderDataloaderConfig, Config, @@ -33,7 +35,7 @@ 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: Dict[str, Config] = {} base_configs["instant_ngp"] = Config( method_name="instant_ngp", trainer=TrainerConfig(mixed_precision=True), diff --git a/nerfactory/configs/utils.py b/nerfactory/configs/utils.py index e2c1567d3..43f20ea06 100644 --- a/nerfactory/configs/utils.py +++ b/nerfactory/configs/utils.py @@ -24,20 +24,14 @@ 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: +def cli_from_base_configs(base_library: Dict[str, T]) -> 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()) @@ -52,8 +46,7 @@ def cli_from_base_configs(base_library: Dict[str, T], eval_mode=False) -> T: 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). + # provided; in this case, it simplifies our CLI but makes it less expressive. avoid_subparsers=True, ) diff --git a/scripts/benchmarking/launch_eval_blender.sh b/scripts/benchmarking/launch_eval_blender.sh index ccf153ca1..065528ea4 100755 --- a/scripts/benchmarking/launch_eval_blender.sh +++ b/scripts/benchmarking/launch_eval_blender.sh @@ -58,10 +58,10 @@ len=${#GPU_IDX[@]} 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} \ - --eval.checkpoint-dir=${checkpoint_dir} \ - --eval.output-filename=${output_dir}/${base_config_name}/blender_${dataset}_${month}-${date}-${year}_${seconds}.json & + config_path="${output_dir}/blender_${dataset}_${month}-${date}-${year}/${base_config_name}/${year}-${month}-${date}_${seconds}/nerfactory_models/" + python scripts/run_eval.py compute-psnr \ + --load-config=${config_path} \ + --output-path=${output_dir}/${base_config_name}/blender_${dataset}_${month}-${date}-${year}_${seconds}.json & echo "Launched ${checkpoint_dir} on gpu ${GPU_IDX[$idx]}" # update gpu @@ -71,4 +71,4 @@ for dataset in ${DATASETS[@]}; do ((idx=idx+1)) fi -done \ No newline at end of file +done diff --git a/scripts/run_eval.py b/scripts/run_eval.py index 205d4dc17..7b1bb2f19 100644 --- a/scripts/run_eval.py +++ b/scripts/run_eval.py @@ -1,16 +1,22 @@ """ run_eval.py """ +from __future__ import annotations + +import dataclasses import json import logging import os -import pickle +from dataclasses import dataclass from pathlib import Path -from typing import Dict +from typing import Dict, Literal, Tuple, Union +import dcargs import mediapy as media import torch +import yaml from tqdm import tqdm +from typing_extensions import assert_never from nerfactory.cameras.camera_paths import ( get_interpolated_camera_path, @@ -18,7 +24,6 @@ from nerfactory.cameras.camera_paths import ( ) from nerfactory.cameras.cameras import Cameras 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 @@ -62,7 +67,7 @@ def _load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path: return load_path -def render_stats_dict(pipeline: Pipeline) -> Dict[str, float]: +def _render_stats_dict(pipeline: Pipeline) -> Dict[str, float]: """Helper function to evaluate the pipeline on a dataloader. Args: @@ -86,7 +91,7 @@ def render_stats_dict(pipeline: Pipeline) -> Dict[str, float]: return {"avg psnr": avg_psnr, "avg rays per sec": avg_rays_per_sec, "avg fps": avg_fps} -def render_trajectory_video( +def _render_trajectory_video( pipeline: Pipeline, cameras: Cameras, output_filename: Path, @@ -121,36 +126,51 @@ def render_trajectory_video( media.write_video(output_filename, images, fps=fps) -def main(config: cfg.Config): - """Evaluate trained model. This evaluation can either render a trajectory or compute the eval psnr. +def _eval_setup(config_path: Path) -> Tuple[cfg.Config, Pipeline, Path]: + """Shared setup for loading a saved pipeline for evaluation. Args: - config_name: Name of the config file to use. - checkpoint_dir: Directory to load the checkpoint from. - rendered_output_name: Name of the renderer output to use (rgb, depth, etc.). - method: Method to use for evaluation. PSNR computes metrics, TRAJ renders a trajectory. - traj: Trajectory to render. - output_filename: Name of the output file. - rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution. - 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: Path to config YAML file. - config.trainer.load_dir = config.eval.checkpoint_dir + Returns: + Loaded config, pipeline module, and corresponding checkpoint. + """ + # load save config + config = yaml.load(config_path.read_text(), Loader=yaml.Loader) + assert isinstance(config, cfg.Config) + + # load checkpoints from wherever they were saved + # TODO: expose the ability to choose an arbitrary checkpoint + config.trainer.load_dir = config.trainer.model_dir config.pipeline.dataloader.eval_image_indices = None - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # setup pipeline (which includes the dataloaders) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") pipeline = config.pipeline.setup(device=device, test_mode=True) + assert isinstance(pipeline, Pipeline) pipeline.eval() - # load checkpointed information - checkpoint_name = _load_checkpoint(config.trainer, pipeline) - if config.eval.method == cfg.MethodType.PSNR: - assert config.eval.output_filename.suffix == ".json" - stats_dict = render_stats_dict(pipeline) + # load checkpointed information + checkpoint_path = _load_checkpoint(config.trainer, pipeline) + + return config, pipeline, checkpoint_path + + +@dataclass +class ComputePSNR: + """Load a checkpoint, compute some PSNR metrics, and save to a JSON.""" + + # Path to config YAML file. + load_config: Path + # Name of the output file. + output_path: Path = Path("output.json") + + def main(self) -> None: + """Main function.""" + config, pipeline, checkpoint_path = _eval_setup(self.load_config) + + assert self.output_path.suffix == ".json" + stats_dict = _render_stats_dict(pipeline) avg_psnr = stats_dict["avg psnr"] avg_rays_per_sec = stats_dict["avg rays per sec"] avg_fps = stats_dict["avg fps"] @@ -158,52 +178,64 @@ def main(config: cfg.Config): print(f"Avg. Rays / Sec: {human_format(avg_rays_per_sec)}") print(f"Avg. FPS: {avg_fps:0.4f}") # save output to some file - config.eval.output_filename.parent.mkdir(parents=True, exist_ok=True) + self.output_path.parent.mkdir(parents=True, exist_ok=True) benchmark_info = { "experiment_name": config.experiment_name, "method_name": config.method_name, - "checkpoint": str(checkpoint_name), + "checkpoint": str(checkpoint_path), "results": stats_dict, } - 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}") + self.output_path.write_text(json.dumps(benchmark_info, indent=2), "utf8") + print(f"Saved results to: {self.output_path}") - 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" + +@dataclasses.dataclass +class RenderTrajectory: + """Load a checkpoint, render a trajectory, and save to a video file.""" + + # Path to config YAML file. + load_config: Path + # Name of the renderer output to use. rgb, depth, etc. + rendered_output_name: str = "rgb" + # Trajectory to render. + traj: Literal["spiral", "interp"] = "spiral" + # Scaling factor to apply to the camera image resolution. + rendered_resolution_scaling_factor: float = 1.0 + # Name of the output file. + output_path: Path = Path("output.mp4") + + def main(self) -> None: + """Main function.""" + config, pipeline, _ = _eval_setup(self.load_config) # TODO(ethan): use camera information from parsing args - if config.eval.traj == cfg.TrajectoryType.SPIRAL: + if self.traj == "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 config.eval.traj == cfg.TrajectoryType.INTERP: + elif self.traj == "interp": cameras = pipeline.dataloader.eval_dataloader.get_camera(image_idx=[0, 10]) camera_path = get_interpolated_camera_path(cameras, steps=30) else: - raise NotImplementedError + assert_never(self.traj) - render_trajectory_video( + _render_trajectory_video( pipeline, camera_path, - 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, + output_filename=self.output_path, + rendered_output_name=self.rendered_output_name, + rendered_resolution_scaling_factor=self.rendered_resolution_scaling_factor, + num_rays_per_chunk=config.pipeline.dataloader.eval_num_rays_per_chunk, ) if __name__ == "__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) + # A Union over dataclass types will create a subcommand for each type. + # + # TODO: it would make sense to split this script up into separate scripts. + # - To reduce duplicate code, some "upstream" refactor could also simplify the + # shared `_load_checkpoint()` and `_eval_setup()` helpers. The high-level + # operations implemented by each seem fairly universal; ideally the checkpoint + # loading logic, for example, would be the same as what's used for loading a + # checkpoint when resuming a training run. + dcargs.cli(Union[ComputePSNR, RenderTrajectory]).main() diff --git a/scripts/run_train.py b/scripts/run_train.py index f7992fd8e..7bd071a4e 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -3,7 +3,6 @@ run_train_nerf.py """ import logging -import pickle import random import socket import traceback @@ -14,6 +13,7 @@ import numpy as np import torch import torch.distributed as dist import torch.multiprocessing as mp +import yaml from nerfactory.configs import base as cfg from nerfactory.configs.utils import cli_from_base_configs @@ -25,7 +25,7 @@ logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=loggin DEFAULT_TIMEOUT = timedelta(minutes=30) # speedup for when input size to model doesn't change (much) -torch.backends.cudnn.benchmark = True +torch.backends.cudnn.benchmark = True # type: ignore def _find_free_port() -> str: @@ -128,6 +128,7 @@ def launch( config (Config, optional): config file specifying training regimen. Defaults to None. timeout (timedelta, optional): timeout of the distributed workers. Defaults to DEFAULT_TIMEOUT. """ + assert config is not None world_size = num_machines * num_gpus_per_machine if world_size == 0: # Using only CPU and one process. @@ -165,6 +166,7 @@ def launch( ), join=False, ) + assert process_context is not None try: process_context.join() except KeyboardInterrupt: @@ -195,19 +197,16 @@ if __name__ == "__main__": 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(f"Loading pre-set config from: {instantiated_config.trainer.load_config}") + instantiated_config = yaml.load(instantiated_config.trainer.load_config.read_text(), Loader=yaml.Loader) + + # print 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)) + logging.info(instantiated_config) + assert instantiated_config.base_dir is not None + instantiated_config.base_dir.mkdir(parents=True, exist_ok=True) + config_yaml_path = instantiated_config.base_dir / "config.yml" + logging.info(f"Saving config to: {config_yaml_path}") + config_yaml_path.write_text(yaml.dump(instantiated_config), "utf8") + main(config=instantiated_config)