diff --git a/.vscode/launch.json b/.vscode/launch.json index 732a0399e..72a0b0534 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -26,7 +26,7 @@ "program": "scripts/run_train.py", "console": "integratedTerminal", "args": [ - "machine_config.num_gpus=0" + "machine.num_gpus=0" ] }, { diff --git a/README.md b/README.md index 522e28dd5..eb01b519c 100644 --- a/README.md +++ b/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 diff --git a/configs/graph_default.yaml b/configs/graph_default.yaml index b87155e30..56f2688b5 100644 --- a/configs/graph_default.yaml +++ b/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 diff --git a/configs/tests/test_dryrun.yaml b/configs/tests/test_dryrun.yaml index 49b7c62fc..389ee8f5a 100644 --- a/configs/tests/test_dryrun.yaml +++ b/configs/tests/test_dryrun.yaml @@ -1,7 +1,7 @@ defaults: - graph_default -machine_config: +machine: num_gpus: 0 trainer: diff --git a/pyrad/data/dataloader.py b/pyrad/data/dataloader.py index 1577d3d8b..3624f0cff 100644 --- a/pyrad/data/dataloader.py +++ b/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 diff --git a/pyrad/engine/trainer.py b/pyrad/engine/trainer.py index f13e1c2ed..ee0901e0a 100644 --- a/pyrad/engine/trainer.py +++ b/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): diff --git a/pyrad/graphs/base.py b/pyrad/graphs/base.py index 458e8fc1e..33df56c7e 100644 --- a/pyrad/graphs/base.py +++ b/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 diff --git a/pyrad/utils/config.py b/pyrad/utils/config.py new file mode 100644 index 000000000..39ba7bc9c --- /dev/null +++ b/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, + ) diff --git a/pyrad/utils/profiler.py b/pyrad/utils/profiler.py index af7c5e5ac..fc409a82c 100644 --- a/pyrad/utils/profiler.py +++ b/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 = {} diff --git a/pyrad/utils/writer.py b/pyrad/utils/writer.py index 8143c4688..ea7ccd001 100644 --- a/pyrad/utils/writer.py +++ b/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"] = {} diff --git a/scripts/run_train.py b/scripts/run_train.py index 791b1ccde..5d9d3c60c 100644 --- a/scripts/run_train.py +++ b/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, )