diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py index b63295c92..b5d2e1243 100644 --- a/mattport/nerf/graph.py +++ b/mattport/nerf/graph.py @@ -3,6 +3,7 @@ The Graph module contains all trainable parameters. """ import importlib from dataclasses import dataclass +from typing import Optional from torch import nn @@ -20,6 +21,7 @@ class Node: name: str children: dict + visited: Optional[bool] = False def __hash__(self): return hash(self.name) @@ -72,6 +74,19 @@ class Graph(nn.Module): roots.add(curr_module) return roots + def topological_sort(self, curr_node: "Node", ordering_stack: list) -> 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 + for child_node in curr_node.children.values(): + if not child_node.visited: + self.topological_sort(child_node, ordering_stack) + ordering_stack.append(curr_node.name) + def get_module_order(self) -> list: """Generates a graph and determines order of module operations using topological sorting @@ -81,10 +96,13 @@ class Graph(nn.Module): Returns: list: ordering of the module names that should be executed """ + roots = self.construct_graph() + ordering_stack = [] + for root in roots: + if not root.visited: + self.topological_sort(root, ordering_stack) - # call construct graph - # function to do topological sort to get ordering - return self.modules_config.keys() + return ordering_stack[::-1] def forward(self, x): """_summary_ diff --git a/tests/nerf/test_graph.py b/tests/nerf/test_graph.py index 9a0d20260..bdeefde58 100644 --- a/tests/nerf/test_graph.py +++ b/tests/nerf/test_graph.py @@ -1,5 +1,5 @@ +import numpy as np import pytest - from mattport.nerf.graph import Graph @@ -135,7 +135,6 @@ def test_construct_graph_repeat(): 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}} - print(test_graph.modules_config["mlp_0"]["inputs"][1]) roots = test_graph.construct_graph() assert len(roots) == 2 @@ -157,3 +156,30 @@ def test_construct_graph_complex(): targets_p3 = [("encoder_0", 3), ("mlp_2", 0)] check_consistency(roots[0], targets_p3) + + +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() + 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() + 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", "inputs": ["x"], "meta_data": {"out_dim": 8}} + order = test_repeated_graph.get_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() + 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')