diff --git a/.vscode/launch.json b/.vscode/launch.json index 9e87d43a5..980221899 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -12,13 +12,20 @@ "justMyCode": true }, { - "name": "Python: run_train", + "name": "Python: run_train (gpu)", + "type": "python", + "request": "launch", + "program": "scripts/run_train.py", + "console": "integratedTerminal" + }, + { + "name": "Python: run_train (cpu)", "type": "python", "request": "launch", "program": "scripts/run_train.py", "console": "integratedTerminal", "args": [ - "machine_config.num_gpus=1" + "machine_config.num_gpus=0" ] } ] 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/data/default_data.yaml b/configs/data/default_data.yaml new file mode 100644 index 000000000..a42faf749 --- /dev/null +++ b/configs/data/default_data.yaml @@ -0,0 +1,16 @@ +# dataset options +dataset: + data_directory: data/blender/lego + dataset_type: blender + downscale_factor: 4 + +# validation options +validation_image_indices: [0, 10] # Which images to evaluate on. + +# 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 diff --git a/configs/default.yaml b/configs/default.yaml index 74496adc8..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 - 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 766028057..9c5ff4c87 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -29,7 +29,15 @@ logging.getLogger("PIL").setLevel(logging.WARNING) class Trainer: """Training class""" - def __init__(self, config: DictConfig, local_rank: int = 0, world_size: int = 1): + def __init__(self, config: DictConfig, local_rank: int = 0, world_size: int = 1, cpu=False): + """_summary_ + + Args: + config (DictConfig): _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. + """ self.config = config self.local_rank = local_rank self.world_size = world_size @@ -43,21 +51,27 @@ class Trainer: self.optimizers = None self.start_step = 0 # logging variables - self.is_main_thread = local_rank % world_size == 0 - 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.is_main_thread = world_size != 0 and local_rank % world_size == 0 + writer = getattr(mattport.utils.writer, self.config.logging.writer.type) + self.writer = writer(self.is_main_thread, save_dir=self.config.logging.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" @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"]) + def collate_fn(self, batch_list): + """TODO(ethan): I need to replace this. + I'm only using this for multiprocess pickle issues for now. + """ + return collate_batch(batch_list, self.config.data.dataloader.num_rays_per_batch, keep_full_image=False) + @profiler.time_function def setup_datasets(self, dataset_inputs_dict: Dict[str, DatasetInputs]): """_summary_""" @@ -65,18 +79,17 @@ class Trainer: image_filenames=dataset_inputs_dict["train"].image_filenames, downscale_factor=dataset_inputs_dict["train"].downscale_factor, ) + self.train_dataset = CollateIterDataset( self.train_image_dataset, - collate_fn=lambda batch_list: collate_batch( - batch_list, self.config.dataloader.num_rays_per_batch, keep_full_image=False - ), - num_samples_to_collate=self.config.dataloader.num_images_to_sample_from, - num_times_to_repeat=self.config.dataloader.num_times_to_repeat_images, + collate_fn=self.collate_fn, + num_samples_to_collate=self.config.data.dataloader.num_images_to_sample_from, + num_times_to_repeat=self.config.data.dataloader.num_times_to_repeat_images, ) self.train_dataloader = DataLoader( self.train_dataset, batch_size=1, - num_workers=self.config.dataloader.num_workers, + num_workers=self.config.data.dataloader.num_workers, collate_fn=collate_batch_size_one, pin_memory=True, ) @@ -93,12 +106,14 @@ class Trainer: dataset_inputs (DatasetInputs): The inputs which will be used to define the camera parameters. """ self.graph = instantiate( - self.config.graph, intrinsics=dataset_inputs.intrinsics, camera_to_world=dataset_inputs.camera_to_world - ).to(f"cuda:{self.local_rank}") + self.config.graph.network, + intrinsics=dataset_inputs.intrinsics, + camera_to_world=dataset_inputs.camera_to_world, + ).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]) @@ -107,7 +122,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 @@ -165,7 +180,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() @@ -179,14 +194,14 @@ 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: - self.test_image(image_idx=0, step=step) - _ = self.test_image(image_idx=10, step=step) if not self.dry_run else None + 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: + for image_idx in self.config.data.validation_image_indices: + self.test_image(image_idx=image_idx, step=step) self.stats.print_stats(i / num_iterations) self.stats.update_time(Stats.TOTAL_TRAIN_TIME, train_start, time(), step=-1) @@ -196,10 +211,14 @@ class Trainer: def train_iteration(self, batch: dict, step: int): """Run one iteration with a batch of inputs.""" # move batch to correct device - ray_indices = batch["indices"].to(f"cuda:{self.local_rank}") + ray_indices = batch["indices"].to(self.device) graph_outputs = self.graph(ray_indices) - batch["pixels"] = batch["pixels"].to(f"cuda:{self.local_rank}") - losses = self.graph.get_losses(batch, graph_outputs) # or self.graph.module + batch["pixels"] = batch["pixels"].to(self.device) + losses = ( + self.graph.module.get_losses(batch, graph_outputs) + if hasattr(self.graph, "module") + else self.graph.get_losses(batch, graph_outputs) + ) loss_sum, loss_dict = self.get_aggregated_loss(losses) self.optimizers.zero_grad_all() loss_sum.backward() @@ -223,7 +242,7 @@ class Trainer: rgb_coarse = [] rgb_fine = [] for i in range(0, num_rays, chunk_size): - ray_indices = all_ray_indices[i : i + chunk_size].to(f"cuda:{self.local_rank}") + ray_indices = all_ray_indices[i : i + chunk_size].to(self.device) graph_outputs = self.graph(ray_indices) rgb_coarse.append(graph_outputs["rgb_coarse"]) rgb_fine.append(graph_outputs["rgb_fine"]) 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/mattport/utils/writer.py b/mattport/utils/writer.py index 69d5f0e41..02eb1eddb 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -31,9 +31,9 @@ def get_tensorboard_name(name: str, group: str = None, prefix: str = None): class Writer: """Writer class""" - def __init__(self, local_rank: int, world_size: int, save_dir: str): + def __init__(self, is_main_thread: bool, save_dir: str): + self.is_main_thread = is_main_thread self.save_dir = save_dir - self.is_main_thread = local_rank % world_size == 0 @abstractmethod def write_image( @@ -83,8 +83,8 @@ class Writer: class TensorboardWriter(Writer): """Tensorboard Writer Class""" - def __init__(self, local_rank: int, world_size: int, save_dir: str): - super().__init__(local_rank, world_size, save_dir) + def __init__(self, is_main_thread: bool, save_dir: str): + super().__init__(is_main_thread, save_dir) if self.is_main_thread: self.tb_writer = SummaryWriter(log_dir=self.save_dir) diff --git a/scripts/run_train.py b/scripts/run_train.py index 715c7ff5b..28c587a15 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -96,7 +96,7 @@ def _distributed_worker( assert num_gpus_per_machine <= torch.cuda.device_count() torch.cuda.set_device(local_rank) - _set_random_seed(1234 + local_rank) + _set_random_seed(1234 + local_rank) # NOTE(ethan): will this work for multiple machines? should we use global_rank? _synchronize(world_size, local_rank) output = main_func(local_rank, world_size, config) @@ -105,7 +105,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: DictConfig, cpu: bool = False) -> Any: """Main training function that sets up and runs the trainer per process Args: @@ -116,7 +116,7 @@ def _train(local_rank: int, world_size: int, config: DictConfig) -> Any: Returns: Any: TODO(): determine the return type """ - trainer = Trainer(config, local_rank, world_size) + trainer = Trainer(config, local_rank, world_size, cpu=cpu) trainer.setup() trainer.train() return 0 @@ -143,7 +143,20 @@ def launch( timeout (datetime, optional): timeout of the distributed workers Defaults to DEFAULT_TIMEOUT. """ world_size = num_machines * num_gpus_per_machine - if world_size > 1: + if world_size == 0: + # Using only CPU and one process. + main_func(local_rank=0, world_size=0, config=config, cpu=True) + elif world_size == 1: + # Using one gpu and one process. + try: + main_func(local_rank=0, world_size=1, config=config) + except KeyboardInterrupt: + pass + finally: + if config.logging.enable_profiler: + profiler.PROFILER.print_profile() + elif world_size > 1: + # Using multiple gpus with multiple processes. if dist_url == "auto": assert num_machines == 1, "dist_url=auto not supported in multi-machine jobs." port = _find_free_port() @@ -169,6 +182,7 @@ def launch( try: process_context.join() except KeyboardInterrupt: + logger = logging.getLogger(__name__) for i, process in enumerate(process_context.processes): if process.is_alive(): logger.info("Terminating process %s", str(i)) @@ -176,15 +190,7 @@ def launch( process.join() logger.info("Process %s finished", str(i)) finally: - if config.logging_configs.enable_profiler: - profiler.PROFILER.print_profile() - else: - try: - main_func(local_rank=0, world_size=1, config=config) - except KeyboardInterrupt: - pass - 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 3978d3462..8e3571895 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 @@ -28,61 +28,64 @@ 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 + + # validation options + validation_image_indices: [0] -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_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: - num_gpus: 1 + num_gpus: 0 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 + 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