diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py index 968522996..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 @@ -14,13 +15,26 @@ class GraphInputs: points: float +@dataclass +class Node: + """Node datastructure for graph composition.""" + + name: str + children: dict + visited: Optional[bool] = False + + def __hash__(self): + return hash(self.name) + + class Graph(nn.Module): """_summary_""" def __init__(self, modules_config: dict) -> None: super().__init__() + self.modules_config = modules_config + # calculate input dimensions based on module dependencies self.modules = {} - self.module_order = self.get_module_order(modules_config) 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": @@ -29,9 +43,52 @@ class Graph(nn.Module): 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 + self.module_order = self.get_module_order() - def get_module_order(self, modules_config: dict) -> list: - """Generates a graph to determine the order in which the modules should execute + def construct_graph(self) -> dict: + """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 + """ + 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={}) + 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}) + processed_modules[input_module] = parent_module + else: + processed_modules[input_module].children[module_name] = curr_module + if input_module == "x": + 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 Args: modules_config (dict): module definitions that make up the network @@ -39,7 +96,13 @@ class Graph(nn.Module): Returns: list: ordering of the module names that should be executed """ - return modules_config.keys() + roots = self.construct_graph() + ordering_stack = [] + for root in roots: + if not root.visited: + self.topological_sort(root, ordering_stack) + + return ordering_stack[::-1] def forward(self, x): """_summary_ @@ -49,4 +112,6 @@ class Graph(nn.Module): """ # for name in self.module_order: # x = self.modules[name](x) + + ## can index into previous results using dictionary raise NotImplementedError diff --git a/tests/nerf/test_graph.py b/tests/nerf/test_graph.py index 5dbf9e468..bdeefde58 100644 --- a/tests/nerf/test_graph.py +++ b/tests/nerf/test_graph.py @@ -1,15 +1,40 @@ +import numpy as np import pytest - from mattport.nerf.graph import Graph def create_basic_graph(): - """Create basic test for creating a graph + """Create basic graph structure + a -> b -> c Returns: dict: module definitions for constructing a graph """ basic_graph = { + "mlp_0": { + "class_name": "MLP", + "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}}, + "mlp_1": { + "class_name": "MLP", + "inputs": ["mlp_0"], + "meta_data": {"in_dim": 4, "out_dim": 5, "num_layers": 6, "layer_width": 7}, + }, + } + return basic_graph + + +def create_repeat_graph(): + """Create graph with repeating nodes + a -> b -> c + a -> + + Returns: + dict: module definitions for constructing a graph + """ + repeat_graph = { "encoder_0": {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}}, "mlp_0": { "class_name": "MLP", @@ -22,22 +47,139 @@ def create_basic_graph(): "meta_data": {"in_dim": 4, "out_dim": 5, "num_layers": 6, "layer_width": 7}, }, } - return basic_graph + return repeat_graph + + +def create_complex_graph(): + """Create graph with varied depths, splits and re-joins + a -> b -> d + -> e + -> c + ------> + + Returns: + dict: module definitions for constructing a graph + """ + complex_graph = { + "encoder_0": {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}}, + "mlp_0": { + "class_name": "MLP", + "inputs": ["encoder_0"], + "meta_data": {"in_dim": 0, "out_dim": 1, "num_layers": 2, "layer_width": 3}, + }, + "mlp_1": { + "class_name": "MLP", + "inputs": ["encoder_0"], + "meta_data": {"in_dim": 4, "out_dim": 5, "num_layers": 6, "layer_width": 7}, + }, + "mlp_2": { + "class_name": "MLP", + "inputs": ["encoder_0", "mlp_0"], + "meta_data": {"in_dim": 4, "out_dim": 5, "num_layers": 6, "layer_width": 7}, + }, + "mlp_3": { + "class_name": "MLP", + "inputs": ["mlp_0"], + "meta_data": {"in_dim": 4, "out_dim": 5, "num_layers": 6, "layer_width": 7}, + }, + } + return complex_graph + + +def check_consistency(root: "Node", targets: list): + """rudimentary check on predicted graph versus expected target ordering""" + curr_pointer = root + for i, (target_name, target_children) in enumerate(targets): + assert curr_pointer.name == target_name + assert len(curr_pointer.children) == target_children + if len(curr_pointer.children) > 0: + assert targets[i + 1][0] in curr_pointer.children + curr_pointer = curr_pointer.children[targets[i + 1][0]] 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['mlp_0']).__name__ == 'MLP' - assert type(test_graph.modules['mlp_1']).__name__ == 'MLP' + assert type(test_graph.modules["encoder_0"]).__name__ == "Encoding" + assert type(test_graph.modules["mlp_0"]).__name__ == "MLP" + assert type(test_graph.modules["mlp_1"]).__name__ == "MLP" 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 + + +def test_construct_graph_basic(): + """test construction of dependency graph: ensure correct ordering""" test_graph = Graph(create_basic_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 - \ No newline at end of file + roots = test_graph.construct_graph() + assert len(roots) == 1 + roots = list(roots) + targets = [("encoder_0", 1), ("mlp_0", 1), ("mlp_1", 0)] + check_consistency(roots[0], targets) + + +def test_construct_graph_repeat(): + """test construction of dependency graph: ensure duplicate nodes are ok""" + test_graph = Graph(create_repeat_graph()) + roots = test_graph.construct_graph() + assert len(roots) == 1 + roots = list(roots) + targets = [("encoder_0", 1), ("mlp_0", 1), ("mlp_1", 0)] + 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}} + roots = test_graph.construct_graph() + assert len(roots) == 2 + + +def test_construct_graph_complex(): + """test construction of dependency graph: ensure different depths are ok""" + test_graph = Graph(create_complex_graph()) + roots = test_graph.construct_graph() + assert len(roots) == 1 + roots = list(roots) + targets_p0 = [("encoder_0", 3), ("mlp_0", 2), ("mlp_3", 0)] + check_consistency(roots[0], targets_p0) + + targets_p1 = [("encoder_0", 3), ("mlp_0", 2), ("mlp_2", 0)] + check_consistency(roots[0], targets_p1) + + targets_p2 = [("encoder_0", 3), ("mlp_1", 0)] + check_consistency(roots[0], targets_p2) + + 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')