updated logic for calculating inputs dynamically; untested code

mattport/nerf/graph.py

@@ -21,7 +21,9 @@ class Node:
 
     name: str
     children: dict
-    visited: Optional[bool] = False
+    parents: dict
+    visited_order: Optional[bool] = False
+    visited_in_dim: Optional[bool] = False
 
     def __hash__(self):
         return hash(self.name)
@@ -33,18 +35,39 @@ class Graph(nn.Module):
     def __init__(self, modules_config: dict) -> None:
         super().__init__()
         self.modules_config = modules_config
-        # calculate input dimensions based on module dependencies
-        self.modules = {}
+        # create graph and get ordering
+        self.roots = self.construct_graph()
+        self.module_order = self.get_module_order()
+
+        # 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 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
+            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"])
-        # generate dependency ordering for module calls
-        self.module_order = self.get_module_order()
+
+        # calculate input dimensions based on module dependencies
+        for root in self.roots:
+            self.set_in_dim(root)
+
+    def set_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()
+            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_order:
+                self.set_in_dim(child_node)
 
     def construct_graph(self) -> dict:
         """Constructs a dependency graph given the module configuration
@@ -59,17 +82,18 @@ class Graph(nn.Module):
         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
+                curr_module.parents[input_module] = parent_module
                 if input_module == "x":
                     roots.add(curr_module)
         return roots
@@ -81,9 +105,9 @@ class Graph(nn.Module):
             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)
 
@@ -96,10 +120,10 @@ 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]