diff --git a/mattport/nerf/field_modules/mlp.py b/mattport/nerf/field_modules/mlp.py index 57c99a900..25deadd01 100644 --- a/mattport/nerf/field_modules/mlp.py +++ b/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. diff --git a/tests/nerf/field_modules/test_field_outputs.py b/tests/nerf/field_modules/test_field_outputs.py index cec54dd90..e7839ccfd 100644 --- a/tests/nerf/field_modules/test_field_outputs.py +++ b/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) diff --git a/tests/nerf/field_modules/test_mlp.py b/tests/nerf/field_modules/test_mlp.py new file mode 100644 index 000000000..c25c29528 --- /dev/null +++ b/tests/nerf/field_modules/test_mlp.py @@ -0,0 +1,31 @@ +""" +MLP Test +""" +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()