Added structured config support (#58)

* added typed config; idealy would convert graph to typed too....

* added graph as config

.vscode/launch.json

@@ -26,7 +26,7 @@
             "program": "scripts/run_train.py",
             "console": "integratedTerminal",
             "args": [
-                "machine_config.num_gpus=0"
+                "machine.num_gpus=0"
             ]
         },
         {

README.md

@@ -70,7 +70,7 @@ With support for [Hydra](https://hydra.cc/), you can run with other configuratio
 
 ```
 # Run with config changes
-python scripts/run_train.py machine_config.num_gpus=1
+python scripts/run_train.py machine.num_gpus=1
 python scripts/run_train.py data.dataset.downscale_factor=1
 
 # Run with different datasets

configs/graph_default.yaml

@@ -3,7 +3,7 @@ defaults:
   - override hydra/hydra_logging: colorlog
   - _self_
 
-machine_config:
+machine:
   seed: 42
   num_gpus: 1
   num_machines: 1

configs/tests/test_dryrun.yaml

@@ -1,7 +1,7 @@
 defaults:
   - graph_default
 
-machine_config:
+machine:
   num_gpus: 0
 
 trainer:

pyrad/data/dataloader.py

@@ -18,7 +18,7 @@ Data loader.
 
 import random
 from typing import Dict, List, Tuple, Union
-from omegaconf import DictConfig, ListConfig
+from omegaconf import ListConfig
 
 from torchtyping import TensorType
 
@@ -30,14 +30,15 @@ from pyrad.data.pixel_sampler import PixelSampler
 from pyrad.data.structs import DatasetInputs
 from pyrad.data.utils import get_dataset_inputs_from_dataset_config
 from pyrad.utils import profiler
+from pyrad.utils.config import DataConfig
 from pyrad.utils.misc import get_dict_to_torch, instantiate_from_dict_config
 
 
 @profiler.time_function
-def setup_dataset_train(config: DictConfig, device: str) -> Tuple[DatasetInputs, "TrainDataloader"]:
+def setup_dataset_train(config: DataConfig, device: str) -> Tuple[DatasetInputs, "TrainDataloader"]:
     """Helper method to load train dataset
     Args:
-        config (DictConfig): Configuration of training dataset.
+        config (DataConfig): Configuration of training dataset.
         device (str): device to load the dataset to
 
     Returns:
@@ -58,10 +59,10 @@ def setup_dataset_train(config: DictConfig, device: str) -> Tuple[DatasetInputs,
 
 
 @profiler.time_function
-def setup_dataset_eval(config: DictConfig, test_mode: bool, device: str) -> Tuple[DatasetInputs, "EvalDataloader"]:
+def setup_dataset_eval(config: DataConfig, test_mode: bool, device: str) -> Tuple[DatasetInputs, "EvalDataloader"]:
     """Helper method to load test or val dataset based on test/train mode
     Args:
-        config (DictConfig): Configuration of training dataset.
+        config (DataConfig): Configuration of training dataset.
         test_mode (bool): specifies whether you are training/testing mode, to load validation/test data
         device (str): device to load the dataset to
 

pyrad/engine/trainer.py

@@ -21,7 +21,6 @@ from typing import Dict
 
 import torch
 import torch.distributed as dist
-from omegaconf import DictConfig
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torchtyping import TensorType
 
@@ -30,6 +29,7 @@ from pyrad.data.dataloader import EvalDataloader, setup_dataset_eval, setup_data
 from pyrad.graphs.base import setup_graph
 from pyrad.optimizers.optimizers import setup_optimizers
 from pyrad.utils import profiler, writer
+from pyrad.utils.config import Config
 from pyrad.utils.decorators import check_main_thread
 from pyrad.utils.writer import EventName, TimeWriter
 
@@ -40,13 +40,13 @@ class Trainer:
     """Training class
 
     Args:
-        config (DictConfig): _description_
+        config (Config): _description_
         local_rank (int, optional): _description_. Defaults to 0.
         world_size (int, optional): _description_. Defaults to 1.
         cpu (bool, optional): Whether or not to use the CPU.
     """
 
-    def __init__(self, config: DictConfig, local_rank: int = 0, world_size: int = 1):
+    def __init__(self, config: Config, local_rank: int = 0, world_size: int = 1):
         self.config = config
         self.local_rank = local_rank
         self.world_size = world_size
@@ -59,7 +59,7 @@ class Trainer:
         self.optimizers = None
         self.start_step = 0
         # logging variables
-        writer.setup_event_writers(config)
+        writer.setup_event_writers(config.logging, max_iter=config.trainer.max_num_iterations)
         profiler.setup_profiler(config.logging)
 
     def setup(self, test_mode=False):

pyrad/graphs/base.py

@@ -30,17 +30,18 @@ from pyrad.cameras.rays import RayBundle
 from pyrad.data.structs import DatasetInputs, SceneBounds
 from pyrad.graphs.modules.ray_generator import RayGenerator
 from pyrad.utils import profiler
+from pyrad.utils.config import GraphConfig
 from pyrad.utils.misc import get_masked_dict, instantiate_from_dict_config, is_not_none
 
 
 @profiler.time_function
-def setup_graph(config: DictConfig, dataset_inputs: DatasetInputs, device: str) -> "Graph":
+def setup_graph(config: GraphConfig, dataset_inputs: DatasetInputs, device: str) -> "Graph":
     """Setup the graph. The dataset inputs should be set with the training data.
 
     Args:
         dataset_inputs (DatasetInputs): The inputs which will be used to define the camera parameters.
     """
-    graph = instantiate_from_dict_config(config, **dataset_inputs.as_dict())
+    graph = instantiate_from_dict_config(DictConfig(config), **dataset_inputs.as_dict())
     graph.to(device)
     return graph
 

pyrad/utils/config.py

@@ -0,0 +1,149 @@
+# 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 MISSING, dataclass
+from typing import Any, Dict, Optional
+from omegaconf import 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
+    # additional optional parameters here
+    resume_train: Optional[ResumeTrainConfig] = None
+
+
+@dataclass
+class DatasetInputsConfig:
+    """Configuration for dataset specification"""
+
+    data_directory: str = MISSING
+    dataset_format: str = MISSING
+    downscale_factor: int = MISSING
+    alpha_color: str = MISSING
+
+
+@dataclass
+class ImageDatasetConfig:
+    """Configuration for type of dataset to instantiate"""
+
+    _target_: str = MISSING
+
+
+@dataclass
+class DataConfig:
+    """Configuration for train/eval datasets"""
+
+    dataset_inputs_train: DatasetInputsConfig = MISSING
+    image_dataset_train: ImageDatasetConfig = MISSING
+    dataloader_train: Dict[str, Any] = MISSING
+    dataset_inputs_eval: DatasetInputsConfig = MISSING
+    image_dataset_eval: ImageDatasetConfig = MISSING
+    dataloader_eval: Dict[str, Any] = MISSING
+    # additional optional parameters here
+    pixel_sampler: Optional[Dict[str, Any]] = None
+
+
+@dataclass
+class GraphConfig:
+    """Configuration for graph instantiation"""
+
+    _target_: str = MISSING
+    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"
+
+
+@dataclass
+class Config:
+    """Full config contents"""
+
+    machine: MachineConfig = MISSING
+    logging: LoggingConfig = MISSING
+    trainer: TrainerConfig = MISSING
+    experiment_name: str = MISSING
+    method_name: str = MISSING
+    data: DataConfig = MISSING
+    graph: GraphConfig = MISSING
+    optimizers: Dict[str, Any] = 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
+    data = DataConfig(**config.data)
+    graph = GraphConfig(**config.graph)
+    optimizers = config.optimizers
+    return Config(
+        machine=machine,
+        logging=logging,
+        trainer=trainer,
+        experiment_name=experiment_name,
+        method_name=method_name,
+        data=data,
+        graph=graph,
+        optimizers=optimizers,
+    )

pyrad/utils/profiler.py

@@ -19,9 +19,8 @@ import logging
 import time
 from typing import Callable
 
-from omegaconf import DictConfig
-
 from pyrad.utils import comms
+from pyrad.utils.config import LoggingConfig
 from pyrad.utils.decorators import check_main_thread, check_profiler_enabled, decorate_all
 
 PROFILER = []
@@ -41,13 +40,13 @@ def time_function(func: Callable) -> Callable:
     return wrapper
 
 
-def flush_profiler(config: DictConfig):
+def flush_profiler(config: LoggingConfig):
     """Method that checks if profiler is enabled before flushing"""
     if config.enable_profiler and PROFILER:
         PROFILER[0].print_profile()
 
 
-def setup_profiler(config: DictConfig):
+def setup_profiler(config: LoggingConfig):
     """Initialization of profilers"""
     if comms.is_main_process():
         PROFILER.append(Profiler(config))
@@ -57,7 +56,7 @@ def setup_profiler(config: DictConfig):
 class Profiler:
     """Profiler class"""
 
-    def __init__(self, config: DictConfig):
+    def __init__(self, config: LoggingConfig):
         self.config = config
         self.profiler_dict = {}
 

pyrad/utils/writer.py

@@ -28,9 +28,10 @@ import imageio
 import numpy as np
 import torch
 import wandb
-from omegaconf import DictConfig, ListConfig
+from omegaconf import ListConfig
 from torch.utils.tensorboard import SummaryWriter
 from torchtyping import TensorType
+from pyrad.utils.config import LoggingConfig
 
 from pyrad.utils.decorators import check_main_thread, decorate_all
 
@@ -134,12 +135,11 @@ def write_out_storage():
 
 
 @check_main_thread
-def setup_event_writers(config: DictConfig) -> None:
+def setup_event_writers(config: LoggingConfig, max_iter: int) -> None:
     """Initialization of all event writers specified in config"""
-    logging_configs = config.logging.writer
-    for writer_type in logging_configs:
+    for writer_type in config.writer:
         writer_class = getattr(sys.modules[__name__], writer_type)
-        writer_config = logging_configs[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:
@@ -147,9 +147,9 @@ def setup_event_writers(config: DictConfig) -> None:
         EVENT_WRITERS.append(curr_writer)
 
     ## configure all the global buffer basic information
-    GLOBAL_BUFFER["max_iter"] = config.trainer.max_num_iterations
-    GLOBAL_BUFFER["max_buffer_size"] = config.logging.max_buffer_size
-    GLOBAL_BUFFER["steps_per_log"] = config.logging.steps_per_log
+    GLOBAL_BUFFER["max_iter"] = max_iter
+    GLOBAL_BUFFER["max_buffer_size"] = config.max_buffer_size
+    GLOBAL_BUFFER["steps_per_log"] = config.steps_per_log
     GLOBAL_BUFFER["events"] = {}
 
 

scripts/run_train.py

@@ -16,11 +16,13 @@ import numpy as np
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
-from omegaconf import DictConfig
 import yaml
+from hydra.core.config_store import ConfigStore
+from omegaconf import DictConfig
 
 from pyrad.engine.trainer import Trainer
 from pyrad.utils import comms, profiler
+from pyrad.utils.config import Config, setup_config
 
 logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG)
 
@@ -53,7 +55,7 @@ def _distributed_worker(
     num_gpus_per_machine: int,
     machine_rank: int,
     dist_url: str,
-    config: DictConfig,
+    config: Config,
     timeout: datetime = DEFAULT_TIMEOUT,
 ) -> Any:
     """Spawned distributed worker that handles the initialization of process group and handles the
@@ -68,7 +70,7 @@ def _distributed_worker(
         dist_url (str): url to connect to for distributed jobs, including protocol
                         e.g. "tcp://127.0.0.1:8686".
                         Can be set to "auto" to automatically select a free port on localhost
-        config (DictConfig): config file specifying trainng regimen
+        config (Config): config file specifying trainng regimen
         timeout (datetime, optional): timeout of the distributed workers Defaults to DEFAULT_TIMEOUT.
 
     Raises:
@@ -102,7 +104,7 @@ def _distributed_worker(
 
     assert num_gpus_per_machine <= torch.cuda.device_count()
     torch.cuda.set_device(local_rank)
-    _set_random_seed(config.machine_config.seed + global_rank)
+    _set_random_seed(config.machine.seed + global_rank)
     comms.synchronize(world_size)
 
     output = main_func(local_rank, world_size, config)
@@ -111,7 +113,7 @@ def _distributed_worker(
     return output
 
 
-def _train(local_rank: int, world_size: int, config: DictConfig) -> Any:
+def _train(local_rank: int, world_size: int, config: Config) -> Any:
     """Main training function that sets up and runs the trainer per process
 
     Args:
@@ -134,7 +136,7 @@ def launch(
     num_machines: int = 1,
     machine_rank: int = 0,
     dist_url: str = "auto",
-    config: DictConfig = None,
+    config: Config = None,
     timeout: datetime = DEFAULT_TIMEOUT,
 ) -> None:
     """Function that spawns muliple processes to call on main_func
@@ -145,17 +147,17 @@ def launch(
         num_machines (int, optional): total number of machines
         machine_rank (int, optional): rank of this machine. Defaults to 0.
         dist_url (str, optional): url to connect to for distributed jobs. Defaults to "auto".
-        config (DictConfig, optional): config file specifying trainng regimen Defaults to None.
+        config (Config, optional): config file specifying trainng regimen Defaults to None.
         timeout (datetime, optional): timeout of the distributed workers Defaults to DEFAULT_TIMEOUT.
     """
     world_size = num_machines * num_gpus_per_machine
     if world_size == 0:
         # Using only CPU and one process.
-        _set_random_seed(config.machine_config.seed)
+        _set_random_seed(config.machine.seed)
         main_func(local_rank=0, world_size=0, config=config)
     elif world_size == 1:
         # Using one gpu and one process.
-        _set_random_seed(config.machine_config.seed)
+        _set_random_seed(config.machine.seed)
         try:
             main_func(local_rank=0, world_size=1, config=config)
         except KeyboardInterrupt:
@@ -200,9 +202,15 @@ def launch(
             profiler.flush_profiler(config.logging)
 
 
+cs = ConfigStore.instance()
+cs.store(name="graph_default", node=Config)
+
+
 @hydra.main(config_path="../configs", config_name="graph_default.yaml")
 def main(config: DictConfig):
     """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, "r", encoding="utf8") as f:
@@ -213,10 +221,10 @@ def main(config: DictConfig):
 
     launch(
         _train,
-        config.machine_config.num_gpus,
-        num_machines=config.machine_config.num_machines,
-        machine_rank=config.machine_config.machine_rank,
-        dist_url=config.machine_config.dist_url,
+        config.machine.num_gpus,
+        num_machines=config.machine.num_machines,
+        machine_rank=config.machine.machine_rank,
+        dist_url=config.machine.dist_url,
         config=config,
     )