Merge branch 'master' of github.com:ethanweber/mattport into tancik/stratified_and_pdf_sampling

.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"
             ]
         }
     ]

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.
 

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

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

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

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

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"])

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

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 = {}
 

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

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)
 

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()
 
 

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