Updated mlp tests

mattport/nerf/field_modules/mlp.py

@@ -17,7 +17,7 @@ class MLP(FieldModule):
         out_dim: int,
         num_layers: int,
         layer_width: int,
-        activation: Optional[nn.Module] = None,
+        activation: Optional[nn.Module] = nn.ReLU(),
         out_activation: Optional[nn.Module] = None,
     ) -> None:
         """Initialize parameters of mulilayer perceptron.

tests/nerf/field_modules/test_field_outputs.py

@@ -1,5 +1,5 @@
 """
-Render Head Tests
+Field output tests
 """
 import pytest
 import torch
@@ -9,7 +9,7 @@ from mattport.nerf.field_modules import field_outputs
 
 
 def test_field_output():
-    """Test rgb render head"""
+    """Test render output"""
     in_dim = 6
     out_dim = 4
     activation = nn.ReLU()
@@ -26,7 +26,7 @@ def test_field_output():
     assert y.shape[-1] == out_dim
 
 
-def test_density_head():
+def test_density_output():
     """Test rgb output"""
     in_dim = 6
     density_head = field_outputs.DensityFieldOutput(in_dim)
@@ -41,7 +41,7 @@ def test_density_head():
     assert y.shape[-1] == 1
 
 
-def test_rgb_head():
+def test_rgb_output():
     """Test rgb output"""
     in_dim = 6
     rgb_head = field_outputs.RGBFieldOutput(in_dim)

tests/nerf/field_modules/test_mlp.py

@@ -0,0 +1,32 @@
+"""
+MLP Test
+"""
+import pytest
+import torch
+from torch import nn
+
+from mattport.nerf.field_modules import MLP
+
+
+def test_mlp():
+    """Test mlp"""
+    in_dim = 6
+    out_dim = 10
+    num_layers = 2
+    layer_width = 32
+    out_activation = nn.ReLU()
+    mlp = MLP(
+        in_dim=in_dim, out_dim=out_dim, num_layers=num_layers, layer_width=layer_width, out_activation=out_activation
+    )
+    assert mlp.get_out_dim() == out_dim
+
+    x = torch.ones((9, in_dim))
+
+    mlp.build_nn_modules()
+    y = mlp(x)
+
+    assert y.shape[-1] == out_dim
+
+
+if __name__ == "__main__":
+    test_mlp()