Merge pull request #13 from ethanweber/modules_structuring

New logic to calculate input dimensions given dependencies

configs/modules.json

@@ -4,15 +4,13 @@
         "inputs": [
             "x"
         ],
-        "meta_data": {
-            "output_dim": 8
-        }
+        "meta_data": {}
     },
-
     "mlp_0": {
         "class": "MLP",
         "inputs": [
-            "encoder_0", "encoder_0"
+            "encoder_0",
+            "encoder_0"
         ],
         "meta_data": {
             "num_layers": 8,
@@ -26,10 +24,14 @@
     },
     "RGB": {
         "name": "rgb_0",
-        "inputs": ["mlp_0"]
+        "inputs": [
+            "mlp_0"
+        ]
     },
     "Density": {
         "name": "rgb_1",
-        "inputs": ["mlp_0"]
+        "inputs": [
+            "mlp_0"
+        ]
     }
 }

mattport/nerf/graph.py

@@ -3,7 +3,7 @@ The Graph module contains all trainable parameters.
 """
 import importlib
 from dataclasses import dataclass
-from typing import Optional
+from typing import Any, Dict, List, Optional, Set
 
 from torch import nn
 
@@ -20,8 +20,10 @@ class Node:
     """Node datastructure for graph composition."""
 
     name: str
-    children: dict
-    visited: Optional[bool] = False
+    children: Dict[str, "Node"]
+    parents: Dict[str, "Node"]
+    visited_order: Optional[bool] = False
+    visited_in_dim: Optional[bool] = False
 
     def __hash__(self):
         return hash(self.name)
@@ -30,64 +32,91 @@ class Node:
 class Graph(nn.Module):
     """_summary_"""
 
-    def __init__(self, modules_config: dict) -> None:
+    def __init__(self, modules_config: Dict[str, Dict[str, Any]]) -> None:
         super().__init__()
         self.modules_config = modules_config
-        # calculate input dimensions based on module dependencies
-        self.modules = {}
-        for module_name, module_dict in modules_config.items():
-            module = getattr(importlib.import_module("mattport.nerf.modules"), module_dict["class_name"])
-            if module_dict["class_name"] != "Encoding":
-                in_dim = 0
-                for inputs in module_dict["inputs"]:
-                    in_dim += modules_config[inputs]["meta_data"]["out_dim"]
-                module_dict["meta_data"]["in_dim"] = in_dim
-            self.modules[module_name] = module(**module_dict["meta_data"])
-        # generate dependency ordering for module calls
+        # create graph and get ordering
+        self.roots = self.construct_graph()
         self.module_order = self.get_module_order()
 
-    def construct_graph(self) -> dict:
+        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"])
+            if "in_dim" not in module_dict["meta_data"]:
+                module_dict["meta_data"]["in_dim"] = 0
+            self.modules[module_name] = module(**module_dict["meta_data"])
+
+        # calculate input dimensions based on module dependencies
+        for root in self.roots:
+            self.get_in_dim(root)
+
+        # instantiate torch.nn members of network
+        for _, module in self.modules.items():
+            module.build_nn_modules()
+
+    def get_in_dim(self, curr_node: Node) -> None:
+        """Dynamically calculates and sets the input dimensions of the modules based on dependency graph
+
+        Args:
+            curr_node (Node): pointer to current node in process
+        """
+        curr_node.visited_in_dim = True
+        if len(curr_node.parents) > 0:
+            in_dim = 0
+            for parent_name in curr_node.parents.keys():
+                in_dim += self.modules[parent_name].get_out_dim()
+            self.modules[curr_node.name].set_in_dim(in_dim)
+            self.modules_config[curr_node.name]["meta_data"]["in_dim"] = in_dim
+
+        for child_node in curr_node.children.values():
+            if not child_node.visited_in_dim:
+                self.get_in_dim(child_node)
+
+    def construct_graph(self) -> Set[Node]:
         """Constructs a dependency graph given the module configuration
 
         Args:
             modules_config (dict): module definitions that make up the network
 
         Returns:
-            list: root nodes for the constructed dependency graph
+            set: all root nodes of the constructed dependency graph
         """
         processed_modules = {}
         roots = set()
         for module_name, module_dict in self.modules_config.items():
             if not module_name in processed_modules:
-                curr_module = Node(name=module_name, children={})
+                curr_module = Node(name=module_name, children={}, parents={})
                 processed_modules[module_name] = curr_module
             else:
                 curr_module = processed_modules[module_name]
             inputs = module_dict["inputs"]
             for input_module in inputs:
                 if not input_module in processed_modules:
-                    parent_module = Node(name=input_module, children={module_name: curr_module})
+                    parent_module = Node(name=input_module, children={module_name: curr_module}, parents={})
                     processed_modules[input_module] = parent_module
                 else:
                     processed_modules[input_module].children[module_name] = curr_module
                 if input_module == "x":
                     roots.add(curr_module)
+                else:
+                    curr_module.parents[input_module] = parent_module
         return roots
 
-    def topological_sort(self, curr_node: "Node", ordering_stack: list) -> None:
+    def topological_sort(self, curr_node: Node, ordering_stack: List[str]) -> None:
         """utility function to sort the call order in the dependency graph
 
         Args:
             curr_node (Node): pointer to current node in process
             ordering_stack (list): cumulative ordering of graph nodes
         """
-        curr_node.visited = True
+        curr_node.visited_order = True
         for child_node in curr_node.children.values():
-            if not child_node.visited:
+            if not child_node.visited_order:
                 self.topological_sort(child_node, ordering_stack)
         ordering_stack.append(curr_node.name)
 
-    def get_module_order(self) -> list:
+    def get_module_order(self) -> List[str]:
         """Generates a graph and determines order of module operations using topological sorting
 
         Args:
@@ -96,12 +125,11 @@ class Graph(nn.Module):
         Returns:
             list: ordering of the module names that should be executed
         """
-        roots = self.construct_graph()
+        roots = self.roots
         ordering_stack = []
         for root in roots:
-            if not root.visited:
+            if not root.visited_order:
                 self.topological_sort(root, ordering_stack)
-
         return ordering_stack[::-1]
 
     def forward(self, x):

mattport/nerf/modules/__init__.py

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

mattport/nerf/modules/encoding.py

@@ -2,11 +2,11 @@
 Encoding functions
 """
 import torch
-from torch import nn
 from torchtyping import TensorType
+from mattport.nerf.modules.module import Module
 
 
-class Encoding(nn.Module):
+class Encoding(Module):
     """Encode an input tensor. Intended to be subclassed"""
 
     def __init__(self, in_dim: int) -> None:
@@ -19,14 +19,6 @@ class Encoding(nn.Module):
             raise ValueError("Input dimension should be greater than zero")
         self.in_dim = in_dim
 
-    def get_out_dim(self) -> int:
-        """Calculates output dimension of encoding.
-
-        Returns:
-            int: output dimension
-        """
-        raise NotImplementedError
-
     def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
         """Encodes an input tensor.
 

mattport/nerf/modules/mlp.py

@@ -5,9 +5,10 @@ from typing import Optional
 
 from torch import nn
 from torchtyping import TensorType
+from mattport.nerf.modules.module import Module
 
 
-class MLP(nn.Module):
+class MLP(Module):
     """Multilayer perceptron"""
 
     def __init__(
@@ -19,7 +20,7 @@ class MLP(nn.Module):
         activation: Optional[nn.Module] = None,
         out_activation: Optional[nn.Module] = None,
     ) -> None:
-        """Initialize mulilayer perceptron.
+        """Initialize parameters of mulilayer perceptron.
 
         Args:
             in_dim (int): Input layer dimension
@@ -31,17 +32,26 @@ class MLP(nn.Module):
         """
         super().__init__()
         self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.num_layers = num_layers
+        self.layer_width = layer_width
+        self.activation = activation
+        self.out_activation = out_activation
+        self.net = None
+
+    def build_model(self) -> None:
+        """Initialize mulilayer perceptron."""
         layers = []
-        for i in range(num_layers):
+        for i in range(self.num_layers):
             if i == 0:
-                layers.append(nn.Linear(in_dim, layer_width))
+                layers.append(nn.Linear(self.in_dim, self.layer_width))
             else:
-                layers.append(nn.Linear(layer_width, layer_width))
-            if activation is not None:
-                layers.append(activation)
-        layers.append(nn.Linear(layer_width, out_dim))
-        if out_activation is not None:
-            layers.append(out_activation)
+                layers.append(nn.Linear(self.layer_width, self.layer_width))
+            if self.activation is not None:
+                layers.append(self.activation)
+        layers.append(nn.Linear(self.layer_width, self.out_dim))
+        if self.out_activation is not None:
+            layers.append(self.out_activation)
 
         self.net = nn.Sequential(*layers)
 

mattport/nerf/modules/module.py

@@ -0,0 +1,62 @@
+"""
+The module baseclass.
+"""
+from abc import abstractmethod
+
+from torch import nn
+from torchtyping import TensorType
+
+
+class Module(nn.Module):
+    """_summary_
+
+    Args:
+        nn (_type_): _description_
+    """
+
+    def __init__(self) -> None:
+        """Default initialization of module"""
+        super().__init__()
+        self.in_dim = 0
+        self.out_dim = 0
+
+    def build_nn_modules(self) -> None:
+        """Function instantiates any torch.nn members within the module.
+        If none exist, do nothing."""
+
+    def set_in_dim(self, in_dim: int) -> None:
+        """Sets input dimension of encoding
+
+        Args:
+            in_dim (int): input dimension
+        """
+        if in_dim <= 0:
+            raise ValueError("Input dimension should be greater than zero")
+        self.in_dim = in_dim
+
+    def get_out_dim(self) -> int:
+        """Calculates output dimension of encoding.
+
+        Returns:
+            int: output dimension
+        """
+        if not hasattr(self, "out_dim"):
+            raise ValueError("Output dimension has not been set")
+        if self.out_dim <= 0:
+            raise ValueError("Output dimension should be greater than zero")
+        return self.out_dim
+
+    @abstractmethod
+    def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
+        """_summary_
+
+        Args:
+            x (TensorType[..., "input_dim"]): Input tensor to be encoded
+
+        Raises:
+            NotImplementedError: _description_
+
+        Returns:
+            TensorType[..., "output_dim"]: A encoded input tensor
+        """
+        raise NotImplementedError

tests/nerf/test_graph.py

@@ -16,7 +16,7 @@ def create_basic_graph():
             "inputs": ["encoder_0"],
             "meta_data": {"in_dim": 0, "out_dim": 1, "num_layers": 2, "layer_width": 3},
         },
-        "encoder_0": {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}},
+        "encoder_0": {"class_name": "ScalingAndOffset", "inputs": ["x"], "meta_data": {"in_dim": 16}},
         "mlp_1": {
             "class_name": "MLP",
             "inputs": ["mlp_0"],
@@ -35,7 +35,7 @@ def create_repeat_graph():
         dict: module definitions for constructing a graph
     """
     repeat_graph = {
-        "encoder_0": {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}},
+        "encoder_0": {"class_name": "ScalingAndOffset", "inputs": ["x"], "meta_data": {"in_dim": 16}},
         "mlp_0": {
             "class_name": "MLP",
             "inputs": ["encoder_0", "encoder_0"],
@@ -61,7 +61,7 @@ def create_complex_graph():
         dict: module definitions for constructing a graph
     """
     complex_graph = {
-        "encoder_0": {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}},
+        "encoder_0": {"class_name": "ScalingAndOffset", "inputs": ["x"], "meta_data": {"in_dim": 16}},
         "mlp_0": {
             "class_name": "MLP",
             "inputs": ["encoder_0"],
@@ -101,7 +101,7 @@ def test_graph_init():
     """test modules dictionary is instantiated properly"""
     test_graph = Graph(create_basic_graph())
     assert test_graph.modules is not None
-    assert type(test_graph.modules["encoder_0"]).__name__ == "Encoding"
+    assert type(test_graph.modules["encoder_0"]).__name__ == "ScalingAndOffset"
     assert type(test_graph.modules["mlp_0"]).__name__ == "MLP"
     assert type(test_graph.modules["mlp_1"]).__name__ == "MLP"
 
@@ -109,7 +109,6 @@ def test_graph_init():
 def test_input_dimension():
     """test module input dimension calculation"""
     test_graph = Graph(create_repeat_graph())
-    assert not hasattr(test_graph.modules["encoder_0"], "in_dim")
     assert test_graph.modules["mlp_0"].in_dim == 16
     assert test_graph.modules["mlp_1"].in_dim == 1
 
@@ -134,7 +133,11 @@ def test_construct_graph_repeat():
     check_consistency(roots[0], targets)
 
     test_graph.modules_config["mlp_0"]["inputs"][1] = "encoder_1"
-    test_graph.modules_config["encoder_1"] = {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}}
+    test_graph.modules_config["encoder_1"] = {
+        "class_name": "ScalingAndOffset",
+        "inputs": ["x"],
+        "meta_data": {"out_dim": 8},
+    }
     roots = test_graph.construct_graph()
     assert len(roots) == 2
 
@@ -162,28 +165,29 @@ def test_ordering():
     """test dependency ordering of constructed graph"""
     test_basic_graph = Graph(create_basic_graph())
     _ = test_basic_graph.construct_graph()
-    order = test_basic_graph.get_module_order()
+    order = test_basic_graph.module_order
     target = ["encoder_0", "mlp_0", "mlp_1"]
     assert np.array_equal(target, order)
 
     test_repeated_graph = Graph(create_repeat_graph())
     _ = test_repeated_graph.construct_graph()
-    order = test_repeated_graph.get_module_order()
+    order = test_repeated_graph.module_order
     target = ["encoder_0", "mlp_0", "mlp_1"]
     assert np.array_equal(target, order)
 
     test_repeated_graph.modules_config["mlp_0"]["inputs"][1] = "encoder_1"
     test_repeated_graph.modules_config["encoder_1"] = {
-        "class_name": "Encoding",
+        "class_name": "ScalingAndOffset",
         "inputs": ["x"],
-        "meta_data": {"out_dim": 8},
+        "meta_data": {"in_dim": 16},
     }
-    order = test_repeated_graph.get_module_order()
+    test_repeated_graph = Graph(test_repeated_graph.modules_config)
+    order = test_repeated_graph.module_order
     assert order.index("mlp_0") > order.index("encoder_0") and order.index("mlp_0") > order.index("encoder_1")
 
     test_complex_graph = Graph(create_complex_graph())
     _ = test_complex_graph.construct_graph()
-    order = test_complex_graph.get_module_order()
+    order = test_complex_graph.module_order
     assert order.index("mlp_2") > order.index("mlp_0")
     assert order.index("mlp_2") > order.index("encoder_0")
     assert order.index("mlp_3") > order.index("encoder_0")