added base class and methods; done logic for building graph

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,8 @@ class Node:
     """Node datastructure for graph composition."""
 
     name: str
-    children: dict
-    parents: dict
+    children: Dict[str, "Node"]
+    parents: Dict[str, "Node"]
     visited_order: Optional[bool] = False
     visited_in_dim: Optional[bool] = False
 
@@ -32,7 +32,7 @@ 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
         # create graph and get ordering
@@ -51,11 +51,11 @@ class Graph(nn.Module):
         for root in self.roots:
             self.get_in_dim(root)
 
-        # instantiate network
+        # instantiate torch.nn members of network
         for _, module in self.modules.items():
-            module.build_model()
+            module.build_nn_modules()
 
-    def get_in_dim(self, curr_node: "Node") -> None:
+    def get_in_dim(self, curr_node: Node) -> None:
         """Dynamically calculates and sets the input dimensions of the modules based on dependency graph
 
         Args:
@@ -73,14 +73,14 @@ class Graph(nn.Module):
             if not child_node.visited_in_dim:
                 self.get_in_dim(child_node)
 
-    def construct_graph(self) -> dict:
+    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()
@@ -103,7 +103,7 @@ class Graph(nn.Module):
                     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:
@@ -116,7 +116,7 @@ class Graph(nn.Module):
                 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:
@@ -130,7 +130,6 @@ class Graph(nn.Module):
         for root in roots:
             if not root.visited_order:
                 self.topological_sort(root, ordering_stack)
-        print(ordering_stack)
         return ordering_stack[::-1]
 
     def forward(self, x):

mattport/nerf/modules/__init__.py

@@ -1,4 +1,4 @@
 """init modules"""
-from .encoding import Encoding
-from .encoding import ScalingAndOffset
+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,21 +19,6 @@ class Encoding(nn.Module):
             raise ValueError("Input dimension should be greater than zero")
         self.in_dim = in_dim
 
-    def build_model(self) -> None:
-        """TODO: create base module class from which this extends"""
-
-    def set_in_dim(self, in_dim: int) -> None:
-        """TODO: create base module class from which this extends"""
-        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__(
@@ -54,14 +55,6 @@ class MLP(nn.Module):
 
         self.net = nn.Sequential(*layers)
 
-    def set_in_dim(self, in_dim: int) -> None:
-        """TODO: create base module class from which this extends"""
-        self.in_dim = in_dim
-
-    def get_out_dim(self) -> int:
-        """TODO: create base module class from which this extends"""
-        return self.out_dim
-
     def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType[..., "out_dim"]:
         """Process input with a multilayer perceptron.
 

mattport/nerf/modules/module.py

@@ -0,0 +1,63 @@
+"""
+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