Config suggestions (#300)

* Config suggestions

* Add some TODO notes

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

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.

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
+wandb

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"""

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),

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,
     )
 

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
+done

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()

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)