merge

configs/config.yaml

@@ -3,7 +3,11 @@ resume_train:
   load_dir: # pretrained model directory
   load_step: # pretrained model's step number
 
-save_step: 10
+max_num_iterations: 1000000
+steps_per_save: 1000
+steps_per_log: 100
+steps_per_test: 5000
+
 model_dir: /tmp/mattport_models/
 
 # dataset options
@@ -34,10 +38,9 @@ param_groups:
       type: SGD
       lr: 1e-4
     scheduler:
-      type: ExponentialLR 
+      type: ExponentialLR
       gamma: 0.1
 
-
 # machine options
 machine_config:
   num_gpus: 4
@@ -48,4 +51,4 @@ machine_config:
 # graph options
 graph:
   # NOTE(ethan): currently setting to nothing b/c we're still designing the config format...
-  _target_: mattport.nerf.graph.nerf.NeRFGraph
+  _target_: mattport.nerf.graph.nerf.NeRFGraph

mattport/nerf/optimizers.py

@@ -47,4 +47,4 @@ class Optimizers:
             param_name (str): name of scheduler to step forward
         """
         if self.config.param_name.scheduler:
-            self.schedulers[param_name].step()
+            self.schedulers[param_name].step()

mattport/nerf/trainer.py

@@ -120,21 +120,26 @@ class Trainer:
 
     def train(self) -> None:
         """_summary_"""
-        num_epochs = 10
-        for _ in range(num_epochs):
-            self.train_iteration()
-
-    def train_iteration(self):
-        """_summary_"""
-        num_iters = 100
-        for step in range(self.start_step, self.start_step + num_iters):
+        num_iterations = 1000
+        for step in range(self.start_step, self.start_step + num_iterations):
             batch = next(iter(self.train_dataloader))
-            # move batch to correct device
-            ray_indices = batch.indices.to(f"cuda:{self.local_rank}")
-            graph_outputs = self.graph(ray_indices)
-            batch.pixels = batch.pixels.to(f"cuda:{self.local_rank}")
-            losses = self.graph.module.get_losses(batch, graph_outputs)
-            self.local_writer.write_text(f"{step}/{num_iters}: losses", losses)
+            losses = self.train_iteration(step, batch)
+            self.local_writer.write_text(f"{step}/{num_iterations}: losses", losses)
             if step % self.config.save_step == 0:
                 self.save_checkpoint(self.config.model_dir, step)
                 self.local_writer.write_text(f"ckpt at {step}", f"{losses}\n")
+
+    def train_iteration(self, step: int, batch: dict):
+        """Run one iteration with a batch of inputs."""
+        # move batch to correct device
+        ray_indices = batch.indices.to(f"cuda:{self.local_rank}")
+        graph_outputs = self.graph(ray_indices)
+        batch.pixels = batch.pixels.to(f"cuda:{self.local_rank}")
+        losses = self.graph.module.get_losses(batch, graph_outputs)
+        return losses
+
+    def test(self, image_idx):
+        """Test a specific image."""
+        with torch.no_grad():
+            ray_indices = None
+            

scripts/run_train_nerf.py

@@ -21,6 +21,9 @@ logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=loggin
 
 DEFAULT_TIMEOUT = timedelta(minutes=30)
 
+# speedup for when input size to model doesn't change (much)
+torch.backends.cudnn.benchmark = True
+
 
 def _find_free_port() -> str:
     """Find free port on open socket"""