Merge branch 'master' of github.com:ethanweber/mattport into eng/distributed_training

mattport/nerf/field_modules/__init__.py

@@ -1,4 +1,4 @@
-"""init modules"""
+"""init field modules"""
 from .encoding import Encoding, ScalingAndOffset
 from .mlp import MLP
-from .module import Module
+from .base import FieldModule

mattport/nerf/field_modules/base.py

@@ -1,5 +1,5 @@
 """
-The module baseclass.
+The field module baseclass.
 """
 from abc import abstractmethod
 
@@ -7,12 +7,8 @@ from torch import nn
 from torchtyping import TensorType
 
 
-class Module(nn.Module):
-    """_summary_
-
-    Args:
-        nn (_type_): _description_
-    """
+class FieldModule(nn.Module):
+    """Field modules that can be combined to store and compute the fields."""
 
     def __init__(self) -> None:
         """Default initialization of module"""

mattport/nerf/field_modules/encoding.py

@@ -3,10 +3,10 @@ Encoding functions
 """
 import torch
 from torchtyping import TensorType
-from mattport.nerf.modules.module import Module
+from mattport.nerf.field_modules.base import FieldModule
 
 
-class Encoding(Module):
+class Encoding(FieldModule):
     """Encode an input tensor. Intended to be subclassed"""
 
     def __init__(self, in_dim: int) -> None:

mattport/nerf/field_modules/field_outputs.py

@@ -6,22 +6,31 @@ from typing import Optional
 from torch import nn
 from torchtyping import TensorType
 
+from mattport.nerf.field_modules.base import FieldModule
 
-class RenderHead(nn.Module):
-    """Base Render head"""
 
-    def __init__(self, in_dim: int, out_dim: int, activation: Optional[nn.Module]) -> None:
+class FieldOutput(FieldModule):
+    """Base field output"""
+
+    def __init__(self, in_dim: int, out_dim: int, activation: Optional[nn.Module] = None) -> None:
         """
-
         Args:
             in_dim (int): input dimension
             out_dim (int): output dimension for renderer
             activation (Optional[nn.Module]): output head activation
         """
         super().__init__()
-        layers = [nn.Linear(in_dim, out_dim)]
-        if activation:
-            layers.append(activation)
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.activation = activation
+
+        self.net = None
+
+    def build_nn_modules(self) -> None:
+        layers = [nn.Linear(self.in_dim, self.out_dim)]
+        if self.activation:
+            layers.append(self.activation)
+        print(layers)
         self.net = nn.Sequential(*layers)
 
     def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType[..., "out_dim"]:
@@ -33,19 +42,20 @@ class RenderHead(nn.Module):
         Returns:
             TensorType[..., "out_dim"]: Render head output
         """
-
+        if not self.net:
+            raise SystemError("Render head network not initialized. build_nn_modules() should be called.")
         return self.net(in_tensor)
 
 
-class DensityHead(RenderHead):
-    """Density head"""
+class DensityFieldOutput(FieldOutput):
+    """Density output"""
 
-    def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Softplus) -> None:
+    def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Softplus()) -> None:
         super().__init__(in_dim, 1, activation)
 
 
-class RGBHead(RenderHead):
-    """RGB head"""
+class RGBFieldOutput(FieldOutput):
+    """RGB output"""
 
-    def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Sigmoid) -> None:
+    def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Sigmoid()) -> None:
         super().__init__(in_dim, 3, activation)

mattport/nerf/field_modules/mlp.py

@@ -5,10 +5,10 @@ from typing import Optional
 
 from torch import nn
 from torchtyping import TensorType
-from mattport.nerf.modules.module import Module
+from mattport.nerf.field_modules.base import FieldModule
 
 
-class MLP(Module):
+class MLP(FieldModule):
     """Multilayer perceptron"""
 
     def __init__(

mattport/nerf/graph.py

@@ -43,7 +43,7 @@ class Graph(nn.Module):
         self.modules = {}
         # initialize graph with known input dimensions; set default in_dim to 0
         for module_name, module_dict in modules_config.items():
-            module = getattr(importlib.import_module("mattport.nerf.modules"), module_dict.class_name)
+            module = getattr(importlib.import_module("mattport.nerf.field_modules"), module_dict.class_name)
             if not module_dict.meta_data.in_dim:
                 module_dict.meta_data.in_dim = 0
             self.modules[module_name] = module(**module_dict.meta_data)

mattport/nerf/loss.py

@@ -0,0 +1,9 @@
+"""
+Collection of Losses
+"""
+from torch import nn
+
+L1Loss = nn.L1Loss
+MSELoss = nn.MSELoss
+
+LOSSES = {"L1": L1Loss, "MSE": MSELoss}

mattport/nerf/renderers.py

@@ -7,13 +7,13 @@ import torch
 from torch import nn
 from torchtyping import TensorType
 
-from mattport.nerf import render_heads as rh
+from mattport.nerf.field_modules import field_outputs
 
 
 class Renderer(nn.Module):
     """Base Renderer. Intended to be subclassed"""
 
-    def required_render_heads(self) -> Set[rh.RenderHead]:
+    def required_render_heads(self) -> Set[field_outputs.FieldOutput]:
         """
         Returns:
             Set[RenderHead]: RenderHeads required for this renderer.
@@ -45,8 +45,8 @@ class RGB(Renderer):
         self.rgb_name = rgb_name
         self.density_name = density_name
 
-    def required_render_heads(self) -> Set[rh.RenderHead]:
-        return set(rh.RGBHead, rh.DensityHead)
+    def required_render_heads(self) -> Set[field_outputs.FieldOutput]:
+        return set(field_outputs.RGBFieldOutput, field_outputs.DensityFieldOutput)
 
     def forward(
         self,
@@ -83,8 +83,8 @@ class Depth(Renderer):
         Renderer (_type_): _description_
     """
 
-    def required_render_heads(self) -> Set[rh.RenderHead]:
-        return set(rh.DensityHead)
+    def required_render_heads(self) -> Set[field_outputs.FieldOutput]:
+        return set(field_outputs.DensityFieldOutput)
 
     def forward(
         self,

tests/nerf/field_modules/test_encodings.py

@@ -3,7 +3,7 @@ Encoding Tests
 """
 import pytest
 import torch
-from mattport.nerf.modules import encoding
+from mattport.nerf.field_modules import encoding
 
 
 def test_scaling_and_offset():

tests/nerf/field_modules/test_field_outputs.py

@@ -0,0 +1,56 @@
+"""
+Render Head Tests
+"""
+import pytest
+import torch
+from torch import nn
+
+from mattport.nerf.field_modules import field_outputs
+
+
+def test_field_output():
+    """Test rgb render head"""
+    in_dim = 6
+    out_dim = 4
+    activation = nn.ReLU()
+    render_head = field_outputs.FieldOutput(in_dim=in_dim, out_dim=out_dim, activation=activation)
+    assert render_head.get_out_dim() == out_dim
+
+    x = torch.ones((9, in_dim))
+
+    with pytest.raises(SystemError):
+        render_head(x)
+    render_head.build_nn_modules()
+    y = render_head(x)
+
+    assert y.shape[-1] == out_dim
+
+
+def test_density_head():
+    """Test rgb output"""
+    in_dim = 6
+    density_head = field_outputs.DensityFieldOutput(in_dim)
+    assert density_head.get_out_dim() == 1
+
+    x = torch.ones((9, in_dim))
+    with pytest.raises(SystemError):
+        density_head(x)
+    density_head.build_nn_modules()
+    y = density_head(x)
+
+    assert y.shape[-1] == 1
+
+
+def test_rgb_head():
+    """Test rgb output"""
+    in_dim = 6
+    rgb_head = field_outputs.RGBFieldOutput(in_dim)
+    assert rgb_head.get_out_dim() == 3
+
+    x = torch.ones((9, in_dim))
+    with pytest.raises(SystemError):
+        rgb_head(x)
+    rgb_head.build_nn_modules()
+    y = rgb_head(x)
+
+    assert y.shape[-1] == 3

tests/nerf/graph/test_graph.py

@@ -1,5 +1,7 @@
+"""
+Tests of graph construction
+"""
 import numpy as np
-import pytest
 from hydra import compose, initialize
 from omegaconf import open_dict
 
@@ -28,7 +30,7 @@ def instantiate_graph(config_name: str) -> Graph:
 
 
 def test_graph_init():
-    """test modules dictionary is instantiated properly"""
+    """test field modules dictionary is instantiated properly"""
     test_graph, _ = instantiate_graph("test_basic_graph")
     assert test_graph.modules is not None
     assert type(test_graph.modules["encoder_0"]).__name__ == "ScalingAndOffset"
@@ -37,7 +39,7 @@ def test_graph_init():
 
 
 def test_input_dimension():
-    """test module input dimension calculation"""
+    """test field module input dimension calculation"""
     test_graph, _ = instantiate_graph("test_repeat_graph")
     assert test_graph.modules["mlp_0"].in_dim == 16
     assert test_graph.modules["mlp_1"].in_dim == 1

tests/nerf/test_encoding.py

@@ -1,9 +0,0 @@
-import pytest
-
-
-def test_positional_encoding():
-    assert True
-
-
-def test_lower_case():
-    assert "MATTPORT".lower() == "mattport"

tests/nerf/test_ray_generator.py

@@ -4,7 +4,7 @@ Test ray generation.
 import torch
 
 from mattport.nerf.dataset.blender_dataset import load_blender_data
-from mattport.nerf.modules.ray_generator import RayGenerator
+from mattport.nerf.field_modules.ray_generator import RayGenerator
 from mattport.structures.rays import RayBundle
 from mattport.utils.io import get_absolute_path