From 8c1e9d7f1ac32407bdd294253e7dc37f2d61130c Mon Sep 17 00:00:00 2001
From: Matt Tancik <tancik@berkeley.edu>
Date: Tue, 10 May 2022 10:18:20 -0700
Subject: [PATCH 1/2] Better align with nerf

---
 configs/graph/default_graph.yaml    | 2 +-
 mattport/nerf/dataset/utils.py      | 2 +-
 mattport/nerf/graph/vanilla_nerf.py | 5 +++--
 mattport/nerf/sampler.py            | 2 +-
 mattport/nerf/trainer.py            | 2 ++
 5 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/configs/graph/default_graph.yaml b/configs/graph/default_graph.yaml
index 3efb85720..819b333e4 100644
--- a/configs/graph/default_graph.yaml
+++ b/configs/graph/default_graph.yaml
@@ -12,7 +12,7 @@ resume_train:
 # 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
+  near_plane: 2.0
   far_plane: 6.0
   num_coarse_samples: 64
   num_importance_samples: 128
diff --git a/mattport/nerf/dataset/utils.py b/mattport/nerf/dataset/utils.py
index 627cd7f93..7cd1bcfd1 100644
--- a/mattport/nerf/dataset/utils.py
+++ b/mattport/nerf/dataset/utils.py
@@ -18,7 +18,7 @@ def get_dataset_inputs_dict(
     dataset_type: str,
     downscale_factor: int = 1,
     alpha_color: Optional[Union[str, list, ListConfig]] = None,
-    splits: Tuple[str] = ("train", "val"),
+    splits: Tuple[str] = ("train", "val", "test"),
 ) -> Dict[str, DatasetInputs]:
     """Returns the dataset inputs, which will be used with an ImageDataset and RayGenerator.
     # TODO: implement the `test` split, which will have depths and normals, etc.
diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py
index 0c67819c7..8d5a62d0d 100644
--- a/mattport/nerf/graph/vanilla_nerf.py
+++ b/mattport/nerf/graph/vanilla_nerf.py
@@ -48,7 +48,8 @@ class NeRFField(nn.Module):
         self.mlp_rgb = MLP(
             in_dim=self.mlp_base.get_out_dim() + self.encoding_dir.get_out_dim(),
             out_dim=self.layer_width // 2,
-            num_layers=1,
+            num_layers=2,
+            layer_width=self.layer_width // 2,
             activation=nn.ReLU(),
         )
         self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_rgb.get_out_dim())
@@ -84,7 +85,7 @@ class NeRFGraph(Graph):
         self,
         intrinsics=None,
         camera_to_world=None,
-        near_plane=1.0,
+        near_plane=2.0,
         far_plane=6.0,
         num_coarse_samples=64,
         num_importance_samples=128,
diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py
index 203c4639e..0b84ac2f6 100644
--- a/mattport/nerf/sampler.py
+++ b/mattport/nerf/sampler.py
@@ -125,7 +125,7 @@ class PDFSampler(nn.Module):
             u = torch.rand(size=(*cdf.shape[:-1], num_samples), device=cdf.device)
         else:
             u = torch.linspace(0.0, 1.0, steps=num_samples, device=cdf.device)
-            u = torch.expand(size=(*cdf.shape[:-1], num_samples))
+            u = u.expand(size=(*cdf.shape[:-1], num_samples))
 
         u = u.contiguous()
         indicies = torch.searchsorted(cdf, u, right=True)
diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py
index e3fce77f5..e1d149d3e 100644
--- a/mattport/nerf/trainer.py
+++ b/mattport/nerf/trainer.py
@@ -240,8 +240,10 @@ class Trainer:
     @profiler.time_function
     def test_image(self, image_idx, step):
         """Test a specific image."""
+        self.graph.eval()
         intrinsics = self.val_image_intrinsics[image_idx]
         camera_to_world = self.val_image_camera_to_world[image_idx]
         outputs = self.graph.get_outputs_for_camera(intrinsics, camera_to_world)
         image = self.val_image_dataset[image_idx]["image"].to(self.device)
         self.graph.log_test_image_outputs(image_idx, step, image, outputs)
+        self.graph.train()

From 2a56fac89c9b05bbe1db44caec9b428025fe894b Mon Sep 17 00:00:00 2001
From: Matt Tancik <tancik@berkeley.edu>
Date: Tue, 10 May 2022 10:22:17 -0700
Subject: [PATCH 2/2] Fix test

---
 mattport/nerf/dataset/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mattport/nerf/dataset/utils.py b/mattport/nerf/dataset/utils.py
index 7cd1bcfd1..627cd7f93 100644
--- a/mattport/nerf/dataset/utils.py
+++ b/mattport/nerf/dataset/utils.py
@@ -18,7 +18,7 @@ def get_dataset_inputs_dict(
     dataset_type: str,
     downscale_factor: int = 1,
     alpha_color: Optional[Union[str, list, ListConfig]] = None,
-    splits: Tuple[str] = ("train", "val", "test"),
+    splits: Tuple[str] = ("train", "val"),
 ) -> Dict[str, DatasetInputs]:
     """Returns the dataset inputs, which will be used with an ImageDataset and RayGenerator.
     # TODO: implement the `test` split, which will have depths and normals, etc.