diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py index b5d2e1243..a59bdbd04 100644 --- a/mattport/nerf/graph.py +++ b/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]