From 6b17ad51c60ab8a4100649038a32487532b012f0 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Fri, 29 Apr 2022 21:29:13 +0000 Subject: [PATCH 1/2] cpu support and multi gpu progress --- .vscode/launch.json | 11 +++++++-- configs/default.yaml | 6 ++--- mattport/nerf/trainer.py | 41 +++++++++++++++++++++++++--------- mattport/utils/writer.py | 8 +++---- scripts/run_train.py | 18 ++++++++++----- tests/configs/test_default.yml | 13 ++++++----- 6 files changed, 65 insertions(+), 32 deletions(-) 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/configs/default.yaml b/configs/default.yaml index dc35f2bf7..03534cf3a 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -12,9 +12,9 @@ logging_configs: # profiler logs run times of functions and prints at end of training enable_profiler: True # writer logs losses and images per iteration - writer: + writer: type: TensorboardWriter - save_dir: './' + save_dir: "./" hydra: run: @@ -45,7 +45,7 @@ dataset: # dataloader options dataloader: - num_workers: 4 + 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 diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 2212c09af..c3d487fb7 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,12 +51,13 @@ class Trainer: self.optimizers = None self.start_step = 0 # logging variables - self.is_main_thread = local_rank % world_size == 0 + self.is_main_thread = world_size != 0 and 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.writer = writer(self.is_main_thread, 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) 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): @@ -57,6 +66,12 @@ class Trainer: 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.dataloader.num_rays_per_batch, keep_full_image=False) + @profiler.time_function def setup_datasets(self, dataset_inputs_dict: Dict[str, DatasetInputs]): """_summary_""" @@ -64,17 +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 - ), + collate_fn=self.collate_fn, num_samples_to_collate=self.config.dataloader.num_images_to_sample_from, num_times_to_repeat=self.config.dataloader.num_times_to_repeat_images, ) self.train_dataloader = DataLoader( self.train_dataset, batch_size=1, + # num_workers=0, num_workers=self.config.dataloader.num_workers, collate_fn=collate_batch_size_one, pin_memory=True, @@ -93,7 +108,7 @@ class Trainer: """ self.graph = instantiate( self.config.graph, intrinsics=dataset_inputs.intrinsics, camera_to_world=dataset_inputs.camera_to_world - ).to(f"cuda:{self.local_rank}") + ).to(self.device) self.setup_optimizers() if self.config.resume_train.load_dir: @@ -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/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 009634f11..49386309c 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -95,7 +95,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) @@ -104,7 +104,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: @@ -115,7 +115,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 @@ -142,7 +142,14 @@ 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. + main_func(local_rank=0, world_size=1, config=config) + 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() @@ -168,14 +175,13 @@ 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)) process.terminate() process.join() logger.info("Process %s finished", str(i)) - else: - main_func(local_rank=0, world_size=1, config=config) @hydra.main(config_path="../configs", config_name="default.yaml") diff --git a/tests/configs/test_default.yml b/tests/configs/test_default.yml index 22932e573..6c286dd9a 100644 --- a/tests/configs/test_default.yml +++ b/tests/configs/test_default.yml @@ -14,9 +14,9 @@ logging_configs: # profiler logs run times of functions and prints at end of training enable_profiler: True # writer logs losses and images per iteration - writer: + writer: type: TensorboardWriter - save_dir: '/tmp/' + save_dir: "/tmp/" hydra: run: @@ -50,7 +50,8 @@ dataloader: num_workers: 4 shuffle: true num_images_to_sample_from: 1 - num_rays_per_batch: 1024 + # num_rays_per_batch: 1024 + num_rays_per_batch: 8 # NOTE(ethan): much smaller than normal! num_times_to_repeat_images: 40 # optimizer options @@ -73,7 +74,7 @@ param_groups: # machine options machine_config: - num_gpus: 1 + num_gpus: 0 num_machines: 1 machine_rank: 0 dist_url: auto @@ -84,5 +85,5 @@ graph: _target_: mattport.nerf.graph.vanilla_nerf.NeRFGraph near_plane: 1.0 far_plane: 6.0 - num_coarse_samples: 64 - num_importance_samples: 64 + num_coarse_samples: 4 # NOTE(ethan): much smaller than normal! + num_importance_samples: 4 # NOTE(ethan): much smaller than normal! From 023e6a3f92085df53850ccba56992d5cfab51a43 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Fri, 29 Apr 2022 22:29:43 +0000 Subject: [PATCH 2/2] merging masters new config --- configs/data/default_data.yaml | 3 +++ mattport/nerf/trainer.py | 19 ++++++++++--------- scripts/run_train.py | 2 +- tests/configs/test_default.yml | 3 +++ 4 files changed, 17 insertions(+), 10 deletions(-) diff --git a/configs/data/default_data.yaml b/configs/data/default_data.yaml index 736eea12e..a42faf749 100644 --- a/configs/data/default_data.yaml +++ b/configs/data/default_data.yaml @@ -4,6 +4,9 @@ dataset: dataset_type: blender downscale_factor: 4 +# validation options +validation_image_indices: [0, 10] # Which images to evaluate on. + # dataloader options dataloader: num_workers: 4 diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 9f547e05f..490e48b9b 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -52,12 +52,11 @@ class Trainer: self.start_step = 0 # logging variables self.is_main_thread = world_size != 0 and local_rank % world_size == 0 - 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) + 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.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 @@ -71,7 +70,7 @@ class Trainer: """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.dataloader.num_rays_per_batch, keep_full_image=False) + 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]): @@ -84,8 +83,8 @@ class Trainer: self.train_dataset = CollateIterDataset( self.train_image_dataset, collate_fn=self.collate_fn, - num_samples_to_collate=self.config.dataloader.num_images_to_sample_from, - num_times_to_repeat=self.config.dataloader.num_times_to_repeat_images, + 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, @@ -107,7 +106,9 @@ 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 + self.config.graph.network, + intrinsics=dataset_inputs.intrinsics, + camera_to_world=dataset_inputs.camera_to_world, ).to(self.device) self.setup_optimizers() @@ -199,8 +200,8 @@ class Trainer: 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 + 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) diff --git a/scripts/run_train.py b/scripts/run_train.py index 3a7c77515..28c587a15 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -153,7 +153,7 @@ def launch( except KeyboardInterrupt: pass finally: - if config.logging_configs.enable_profiler: + if config.logging.enable_profiler: profiler.PROFILER.print_profile() elif world_size > 1: # Using multiple gpus with multiple processes. diff --git a/tests/configs/test_default.yml b/tests/configs/test_default.yml index eba2cf917..8e3571895 100644 --- a/tests/configs/test_default.yml +++ b/tests/configs/test_default.yml @@ -34,6 +34,9 @@ data: data_directory: data/blender/lego_test dataset_type: blender downscale_factor: 4 + + # validation options + validation_image_indices: [0] # dataloader options dataloader: