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/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/configs/default.yaml b/configs/default.yaml index cd23f0aca..2e579dda2 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -20,4 +20,4 @@ machine_config: num_gpus: 1 num_machines: 1 machine_rank: 0 - dist_url: auto + dist_url: auto \ No newline at end of file diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 5a1f4b92c..490e48b9b 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,13 +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.writer.type) - self.writer = writer(local_rank, world_size, save_dir=self.config.logging.writer.save_dir) + 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 def setup(self): @@ -58,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.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,11 +79,10 @@ 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.data.dataloader.num_rays_per_batch, keep_full_image=False - ), + 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, ) @@ -96,7 +109,7 @@ class Trainer: self.config.graph.network, 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.graph.resume_train.load_dir: @@ -187,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) @@ -198,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() @@ -225,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 e59ff6f59..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)) @@ -178,14 +192,6 @@ def launch( finally: if config.logging.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.enable_profiler: - profiler.PROFILER.print_profile() @hydra.main(config_path="../configs", config_name="default.yaml") diff --git a/tests/configs/test_default.yml b/tests/configs/test_default.yml index f515cb984..8e3571895 100644 --- a/tests/configs/test_default.yml +++ b/tests/configs/test_default.yml @@ -34,18 +34,21 @@ data: data_directory: data/blender/lego_test dataset_type: blender downscale_factor: 4 + + # validation options + validation_image_indices: [0] # dataloader options dataloader: num_workers: 4 shuffle: true num_images_to_sample_from: 1 - num_rays_per_batch: 1024 + 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 @@ -67,8 +70,8 @@ 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 + num_importance_samples: 4 # optimizer options param_groups: fields: @@ -85,4 +88,4 @@ graph: lr: 1e-4 scheduler: type: ExponentialLR - gamma: 0.1 + gamma: 0.1 \ No newline at end of file