diff --git a/README.md b/README.md index 0a14f9a73..666ec8db1 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,7 @@ python scripts/run_train.py machine_config.num_gpus=1 ``` # Logging and profiling features -We provide many logging functionalities for timing and/or tracking losses during training. All of these loggers are configurable via `configs/logging_configs.yml` +We provide many logging functionalities for timing and/or tracking losses during training. All of these loggers are configurable via `configs/logging.yml` 1. **Writer**: Logs losses and generated images during training to a specified output stream. Specify the type of writer (Tensorboard, local directory, Weights and Biases), and how often to log in the config. diff --git a/configs/default.yaml b/configs/default.yaml index 0467b522d..2e579dda2 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -1,22 +1,11 @@ defaults: # This makes the logging module logs colorful + - data: default_data.yaml + - graph: default_graph.yaml + - logging: default_logging.yaml - override hydra/job_logging: colorlog - override hydra/hydra_logging: colorlog -logging_configs: - steps_per_log: 10 - # stats tracker logs configurable training stats per iteration to terminal (e.g. rays/s, PSNR) - enable_stats: True - stats_tracker: - max_history: 10 - stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME] # see mattport/utils/stats_tracker.py for options - # writer logs losses and images per iteration - writer: - type: TensorboardWriter - save_dir: "./" - # profiler logs run times of functions and prints at end of training - enable_profiler: True - hydra: run: # Configure output dir of each experiment programmatically from the arguments @@ -26,61 +15,9 @@ hydra: experiment_name: blender_lego method_name: vanilla_nerf -# pretrained loading options -resume_train: - load_dir: # pretrained model directory - load_step: # pretrained model's step number - -max_num_iterations: 1000000 -steps_per_save: null -steps_per_test: 100 - -model_dir: /tmp/mattport_models/ - -# dataset options -dataset: - data_directory: data/blender/lego - dataset_type: blender - downscale_factor: 4 - -# dataloader options -dataloader: - num_workers: 4 # Need to set to 0 if using multiple GPUs (until we fix a bug). - shuffle: true - num_images_to_sample_from: 8 - num_rays_per_batch: 1024 - num_times_to_repeat_images: 40 - -# optimizer options -param_groups: - fields: - optimizer: - type: Adam - lr: 1e-4 - scheduler: - type: StepLR - step_size: 30 # Period of learning rate decay. - gamma: 0.1 # Multiplicative factor of learning rate decay. Default: 0.1. - cameras: - optimizer: - type: SGD - lr: 1e-4 - scheduler: - type: ExponentialLR - gamma: 0.1 - # machine options machine_config: num_gpus: 1 num_machines: 1 machine_rank: 0 - dist_url: auto - -# graph options -graph: - # NOTE(ethan): currently setting to nothing b/c we're still designing the config format... - _target_: mattport.nerf.graph.vanilla_nerf.NeRFGraph - near_plane: 1.0 - far_plane: 6.0 - num_coarse_samples: 64 - num_importance_samples: 64 + dist_url: auto \ No newline at end of file diff --git a/configs/graph/default_graph.yaml b/configs/graph/default_graph.yaml new file mode 100644 index 000000000..9fd43adf5 --- /dev/null +++ b/configs/graph/default_graph.yaml @@ -0,0 +1,36 @@ +# saving model specifications +model_dir: /tmp/mattport_models/ +steps_per_save: null +steps_per_test: 100 +max_num_iterations: 1000000 + +# pretrained loading options +resume_train: + load_dir: # pretrained model directory + load_step: # pretrained model's step number + +# NOTE(ethan): currently setting to nothing b/c we're still designing the config format... +network: + _target_: mattport.nerf.graph.vanilla_nerf.NeRFGraph + near_plane: 1.0 + far_plane: 6.0 + num_coarse_samples: 64 + num_importance_samples: 64 + +# optimizer options for network +param_groups: + fields: + optimizer: + type: Adam + lr: 1e-4 + scheduler: + type: StepLR + step_size: 30 # Period of learning rate decay. + gamma: 0.1 # Multiplicative factor of learning rate decay. Default: 0.1. + cameras: + optimizer: + type: SGD + lr: 1e-4 + scheduler: + type: ExponentialLR + gamma: 0.1 \ No newline at end of file diff --git a/configs/logging/default_logging.yaml b/configs/logging/default_logging.yaml new file mode 100644 index 000000000..fd5e88371 --- /dev/null +++ b/configs/logging/default_logging.yaml @@ -0,0 +1,12 @@ +steps_per_log: 10 +# stats tracker logs configurable training stats per iteration to terminal (e.g. rays/s, PSNR) +enable_stats: True +stats_tracker: + max_history: 10 + stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME] # see mattport/utils/stats_tracker.py for options +# writer logs losses and images per iteration +writer: + type: TensorboardWriter + save_dir: './' +# profiler logs run times of functions and prints at end of training +enable_profiler: True \ No newline at end of file diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 717273087..9f547e05f 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -55,7 +55,7 @@ class Trainer: writer = getattr(mattport.utils.writer, self.config.logging_configs.writer.type) self.writer = writer(self.is_main_thread, save_dir=self.config.logging_configs.writer.save_dir) self.stats = StatsTracker(config, self.is_main_thread) - if not profiler.PROFILER and self.config.logging_configs.enable_profiler: + if not profiler.PROFILER and self.config.logging.enable_profiler: profiler.PROFILER = profiler.Profiler(config, self.is_main_thread) self.dry_run = self.config.get("dry_run", False) self.device = f"cuda:{self.local_rank}" if not cpu else "cpu" @@ -63,7 +63,7 @@ class Trainer: @profiler.time_function def setup(self): """Setup the Trainer by calling other setup functions.""" - dataset_inputs_dict = get_dataset_inputs_dict(**self.config.dataset) + dataset_inputs_dict = get_dataset_inputs_dict(**self.config.data.dataset) self.setup_datasets(dataset_inputs_dict) self.setup_graph(dataset_inputs_dict["train"]) @@ -90,8 +90,7 @@ class Trainer: self.train_dataloader = DataLoader( self.train_dataset, batch_size=1, - # num_workers=0, - num_workers=self.config.dataloader.num_workers, + num_workers=self.config.data.dataloader.num_workers, collate_fn=collate_batch_size_one, pin_memory=True, ) @@ -112,8 +111,8 @@ class Trainer: ).to(self.device) self.setup_optimizers() - if self.config.resume_train.load_dir: - self.load_checkpoint(self.config.resume_train) + if self.config.graph.resume_train.load_dir: + self.load_checkpoint(self.config.graph.resume_train) if self.world_size > 1: self.graph = DDP(self.graph, device_ids=[self.local_rank]) @@ -122,7 +121,7 @@ class Trainer: def setup_optimizers(self): """_summary_""" # TODO(ethan): handle different world sizes - self.optimizers = Optimizers(self.config.param_groups, self.graph.get_param_groups()) + self.optimizers = Optimizers(self.config.graph.param_groups, self.graph.get_param_groups()) def load_checkpoint(self, load_config: DictConfig) -> int: """Load the checkpoint from the given path @@ -180,7 +179,7 @@ class Trainer: def train(self) -> None: """_summary_""" train_start = time() - num_iterations = self.config.max_num_iterations + num_iterations = self.config.graph.max_num_iterations iter_dataset = iter(self.train_dataloader) for i, step in enumerate(range(self.start_step, self.start_step + num_iterations)): data_start = time() @@ -194,12 +193,12 @@ class Trainer: ) self.stats.update_time(Stats.ITER_TRAIN_TIME, iter_start, time(), step=step) - if step != 0 and step % self.config.logging_configs.steps_per_log == 0: + if step != 0 and step % self.config.logging.steps_per_log == 0: self.writer.write_scalar_dict(loss_dict, step, group="Loss", prefix="train-") # TODO: add the learning rates to tensorboard/logging - if step != 0 and self.config.steps_per_save and step % self.config.steps_per_save == 0: - self.save_checkpoint(self.config.model_dir, step) - if step % self.config.steps_per_test == 0: + if step != 0 and self.config.graph.steps_per_save and step % self.config.graph.steps_per_save == 0: + self.save_checkpoint(self.config.graph.model_dir, step) + if step % self.config.graph.steps_per_test == 0: self.test_image(image_idx=0, step=step) _ = self.test_image(image_idx=10, step=step) if not self.dry_run else None self.stats.print_stats(i / num_iterations) diff --git a/mattport/utils/decorators.py b/mattport/utils/decorators.py index 779b1a3df..988ae7c56 100644 --- a/mattport/utils/decorators.py +++ b/mattport/utils/decorators.py @@ -22,7 +22,7 @@ def check_stats_enabled(func: Callable) -> Callable: def wrapper(self, *args, **kwargs): ret = None - if self.config.logging_configs.enable_stats: + if self.config.logging.enable_stats: ret = func(self, *args, **kwargs) return ret @@ -34,7 +34,7 @@ def check_profiler_enabled(func: Callable) -> Callable: def wrapper(self, *args, **kwargs): ret = None - if self.config.logging_configs.enable_stats: + if self.config.logging.enable_stats: ret = func(self, *args, **kwargs) return ret @@ -59,9 +59,9 @@ def check_print_stats_step(func: Callable) -> Callable: def wrapper(self, *args, **kwargs): ret = None if ( - self.step % self.config.logging_configs.steps_per_log == 0 - or (self.config.steps_per_save and self.step % self.config.steps_per_save == 0) - or self.step % self.config.steps_per_test == 0 + self.step % self.config.logging.steps_per_log == 0 + or (self.config.graph.steps_per_save and self.step % self.config.graph.steps_per_save == 0) + or self.step % self.config.graph.steps_per_test == 0 ): ret = func(self, *args, **kwargs) return ret diff --git a/mattport/utils/profiler.py b/mattport/utils/profiler.py index 8f193340b..0d7f12619 100644 --- a/mattport/utils/profiler.py +++ b/mattport/utils/profiler.py @@ -37,7 +37,7 @@ class Profiler: def __init__(self, config: DictConfig, is_main_thread: bool): self.config = config - if self.config.logging_configs.enable_profiler: + if self.config.logging.enable_profiler: self.is_main_thread = is_main_thread self.profiler_dict = {} diff --git a/mattport/utils/stats_tracker.py b/mattport/utils/stats_tracker.py index 119993326..03461c050 100644 --- a/mattport/utils/stats_tracker.py +++ b/mattport/utils/stats_tracker.py @@ -30,15 +30,15 @@ class StatsTracker: def __init__(self, config: DictConfig, is_main_thread: bool): self.config = config - if self.config.logging_configs.enable_stats: + if self.config.logging.enable_stats: self.is_main_thread = is_main_thread - self.max_history = self.config.logging_configs.stats_tracker.max_history + self.max_history = self.config.logging.stats_tracker.max_history self.step = 0 self.stats_dict = {} self.past_stats = [] self.new_key = False self.stats_to_track = set() - for name in self.config.logging_configs.stats_tracker.stats_to_track: + for name in self.config.logging.stats_tracker.stats_to_track: self.stats_to_track.add(Stats[name]) logging.info("Successfully set up StatsTracker") else: @@ -85,7 +85,7 @@ class StatsTracker: if name == Stats.ITER_TRAIN_TIME and Stats.ETA in self.stats_to_track: # update ETA if logging iteration train time - remain_iter = self.config.max_num_iterations - step + remain_iter = self.config.graph.max_num_iterations - step self.stats_dict[Stats.ETA] = remain_iter * self.stats_dict[name] @check_print_stats_step diff --git a/scripts/run_train.py b/scripts/run_train.py index 9ff59a5be..3a7c77515 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -190,7 +190,7 @@ def launch( process.join() logger.info("Process %s finished", str(i)) finally: - if config.logging_configs.enable_profiler: + if config.logging.enable_profiler: profiler.PROFILER.print_profile() diff --git a/tests/configs/test_default.yml b/tests/configs/test_default.yml index 699602110..eba2cf917 100644 --- a/tests/configs/test_default.yml +++ b/tests/configs/test_default.yml @@ -5,7 +5,7 @@ defaults: - override hydra/job_logging: colorlog - override hydra/hydra_logging: colorlog -logging_configs: +logging: steps_per_log: 10 # stats tracker logs configurable training stats per iteration to terminal (e.g. rays/s, PSNR) enable_stats: True @@ -13,9 +13,9 @@ logging_configs: max_history: 10 stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME] # see mattport/utils/stats_tracker.py for options # writer logs losses and images per iteration - writer: + writer: type: TensorboardWriter - save_dir: "/tmp/" + save_dir: './' # profiler logs run times of functions and prints at end of training enable_profiler: True @@ -28,49 +28,20 @@ hydra: experiment_name: blender_lego method_name: vanilla_nerf -# pretrained loading options -resume_train: - load_dir: # pretrained model directory - load_step: # pretrained model's step number +data: + # dataset options + dataset: + data_directory: data/blender/lego_test + dataset_type: blender + downscale_factor: 4 -max_num_iterations: 10 -steps_per_save: null -steps_per_test: 100 - -model_dir: /tmp/mattport_models/ - -# dataset options -dataset: - data_directory: data/blender/lego_test - dataset_type: blender - downscale_factor: 4 - -# dataloader options -dataloader: - num_workers: 4 - shuffle: true - num_images_to_sample_from: 1 - # num_rays_per_batch: 1024 - num_rays_per_batch: 8 # NOTE(ethan): much smaller than normal! - num_times_to_repeat_images: 40 - -# optimizer options -param_groups: - fields: - optimizer: - type: Adam - lr: 1e-4 - scheduler: - type: StepLR - step_size: 30 # Period of learning rate decay. - gamma: 0.1 # Multiplicative factor of learning rate decay. Default: 0.1. - cameras: - optimizer: - type: SGD - lr: 1e-4 - scheduler: - type: ExponentialLR - gamma: 0.1 + # dataloader options + dataloader: + num_workers: 4 + shuffle: true + num_images_to_sample_from: 1 + num_rays_per_batch: 8 + num_times_to_repeat_images: 40 # machine options machine_config: @@ -81,9 +52,37 @@ machine_config: # graph options graph: - # NOTE(ethan): currently setting to nothing b/c we're still designing the config format... - _target_: mattport.nerf.graph.vanilla_nerf.NeRFGraph - near_plane: 1.0 - far_plane: 6.0 - num_coarse_samples: 4 # NOTE(ethan): much smaller than normal! - num_importance_samples: 4 # NOTE(ethan): much smaller than normal! + model_dir: /tmp/mattport_models/ + max_num_iterations: 10 + steps_per_save: null + steps_per_test: 100 + + # pretrained loading options + resume_train: + load_dir: # pretrained model directory + load_step: # pretrained model's step number + + network: + # NOTE(ethan): currently setting to nothing b/c we're still designing the config format... + _target_: mattport.nerf.graph.vanilla_nerf.NeRFGraph + near_plane: 1.0 + far_plane: 6.0 + num_coarse_samples: 4 + num_importance_samples: 4 + # optimizer options + param_groups: + fields: + optimizer: + type: Adam + lr: 1e-4 + scheduler: + type: StepLR + step_size: 30 # Period of learning rate decay. + gamma: 0.1 # Multiplicative factor of learning rate decay. Default: 0.1. + cameras: + optimizer: + type: SGD + lr: 1e-4 + scheduler: + type: ExponentialLR + gamma: 0.1 \ No newline at end of file