diff --git a/configs/default.yaml b/configs/default.yaml index e56cae300..dc35f2bf7 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -3,12 +3,18 @@ defaults: - override hydra/job_logging: colorlog - override hydra/hydra_logging: colorlog -debug: +logging_configs: + # 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 + # profiler logs run times of functions and prints at end of training enable_profiler: True + # writer logs losses and images per iteration + writer: + type: TensorboardWriter + save_dir: './' hydra: run: diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 3ea064bb7..cd8ceabfe 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -13,6 +13,7 @@ from omegaconf import DictConfig from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader +import mattport.utils.writer from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch from mattport.nerf.dataset.utils import get_dataset_inputs @@ -21,7 +22,6 @@ from mattport.nerf.optimizers import Optimizers from mattport.utils import profiler from mattport.utils.decorators import check_main_thread from mattport.utils.stats_tracker import Stats, StatsTracker -from mattport.utils.writer import TensorboardWriter logging.getLogger("PIL").setLevel(logging.WARNING) @@ -45,9 +45,8 @@ class Trainer: self.start_step = 0 # logging variables self.is_main_thread = local_rank % world_size == 0 - self.tensorboard_writer = TensorboardWriter( - local_rank, world_size, save_dir=os.path.join(os.getcwd(), "writer") - ) + writer = getattr(mattport.utils.writer, self.config.logging_configs.writer.type) + self.writer = writer(local_rank, world_size, save_dir=self.config.logging_configs.writer.save_dir) self.stats = StatsTracker(config, self.is_main_thread) profiler.PROFILER = profiler.Profiler(config, self.is_main_thread) @@ -167,7 +166,7 @@ class Trainer: self.stats.update_time(Stats.ITER_TRAIN_TIME, iter_start, time(), step=step) if step != 0 and step % self.config.steps_per_log == 0: - self.tensorboard_writer.write_scalar_dict(loss_dict, step, group="Loss", prefix="train-") + 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) @@ -220,10 +219,8 @@ class Trainer: rgb_fine = torch.cat(rgb_fine).view(image_height, image_width, 3).detach().cpu() combined_image = torch.cat([image, rgb_coarse, rgb_fine], dim=1) - self.tensorboard_writer.write_image( - f"image_idx_{image_idx}-rgb_coarse_fine", combined_image, step, group="test" - ) + self.writer.write_image(f"image_idx_{image_idx}-rgb_coarse_fine", combined_image, step, group="test") fine_psnr = get_psnr(image, rgb_fine) self.stats.update_value(Stats.CURR_TEST_PSNR, fine_psnr, step) - self.tensorboard_writer.write_scalar(f"image_idx_{image_idx}-fine_psnr", fine_psnr, step, group="test") + self.writer.write_scalar(f"image_idx_{image_idx}-fine_psnr", fine_psnr, step, group="test") diff --git a/mattport/utils/decorators.py b/mattport/utils/decorators.py index d76aaf684..359b63ef8 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.debug.enable_stats: + if self.config.logging_configs.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.debug.enable_stats: + if self.config.logging_configs.enable_stats: ret = func(self, *args, **kwargs) return ret diff --git a/mattport/utils/profiler.py b/mattport/utils/profiler.py index c23505233..9966e03ce 100644 --- a/mattport/utils/profiler.py +++ b/mattport/utils/profiler.py @@ -36,7 +36,7 @@ class Profiler: def __init__(self, config: DictConfig, is_main_thread: bool): self.config = config - if self.config.debug.enable_profiler: + if self.config.logging_configs.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 909d39c53..13d8e38a7 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.debug.enable_stats: + if self.config.logging_configs.enable_stats: self.is_main_thread = is_main_thread - self.max_history = self.config.debug.stats_tracker.max_history + self.max_history = self.config.logging_configs.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.debug.stats_tracker.stats_to_track: + for name in self.config.logging_configs.stats_tracker.stats_to_track: self.stats_to_track.add(Stats[name]) logging.info("Successfully set up StatsTracker") else: diff --git a/tests/configs/test_default.yml b/tests/configs/test_default.yml new file mode 100644 index 000000000..7b30f9e57 --- /dev/null +++ b/tests/configs/test_default.yml @@ -0,0 +1,86 @@ +defaults: + # This makes the logging module logs colorful + - override hydra/job_logging: colorlog + - override hydra/hydra_logging: colorlog + +logging_configs: + # 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 + # profiler logs run times of functions and prints at end of training + enable_profiler: True + # writer logs losses and images per iteration + writer: + type: TensorboardWriter + save_dir: '/tmp/' + +hydra: + run: + # Configure output dir of each experiment programmatically from the arguments + # Example "outputs/mnist/classifier/baseline/2021-03-10-141516" + dir: outputs/${experiment_name}/${method_name}/${now:%Y-%m-%d_%H%M%S} + +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: 10 +steps_per_save: null +steps_per_log: 10 +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 + 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 diff --git a/tests/test_run_train.py b/tests/test_run_train.py new file mode 100644 index 000000000..bd6cc24b3 --- /dev/null +++ b/tests/test_run_train.py @@ -0,0 +1,16 @@ +""" +Default test to make sure train runs +""" +import pytest +import omegaconf +from mattport.utils.io import get_absolute_path + +from scripts.run_train import main + + +@pytest.mark.filterwarnings("ignore::DeprecationWarning") +def test_run_train(): + """test run train script works properly""" + test_config = get_absolute_path("./tests/configs/test_default.yml") + cfg = omegaconf.OmegaConf.load(test_config) + main(cfg)