From b6f878f227cde55dd049f44e38bed6906b65bfb9 Mon Sep 17 00:00:00 2001
From: Ethan Weber <eweb0124@gmail.com>
Date: Thu, 21 Apr 2022 20:29:43 +0000
Subject: [PATCH] merge

---
 configs/config.yaml         | 11 +++++++----
 mattport/nerf/optimizers.py |  2 +-
 mattport/nerf/trainer.py    | 33 +++++++++++++++++++--------------
 scripts/run_train_nerf.py   |  3 +++
 4 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/configs/config.yaml b/configs/config.yaml
index 41323f313..397a54279 100644
--- a/configs/config.yaml
+++ b/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
\ No newline at end of file
+  _target_: mattport.nerf.graph.nerf.NeRFGraph
diff --git a/mattport/nerf/optimizers.py b/mattport/nerf/optimizers.py
index e88b8e442..6e59836ca 100644
--- a/mattport/nerf/optimizers.py
+++ b/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()
\ No newline at end of file
diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py
index 3859e2a38..ab5503dc0 100644
--- a/mattport/nerf/trainer.py
+++ b/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
+            
\ No newline at end of file
diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py
index b7f0835c4..abe58c7ae 100644
--- a/scripts/run_train_nerf.py
+++ b/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"""