From 444c387bee3523a897d6622a1aa46db6dbbd4be4 Mon Sep 17 00:00:00 2001 From: evonneng Date: Sat, 9 Apr 2022 19:45:11 +0000 Subject: [PATCH] done logic constructing graph; validated with weaker testing of construction --- mattport/nerf/graph.py | 54 +++++++++++++++-- tests/nerf/test_graph.py | 125 ++++++++++++++++++++++++++++++++++++--- 2 files changed, 166 insertions(+), 13 deletions(-) diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py index 968522996..389ed6867 100644 --- a/mattport/nerf/graph.py +++ b/mattport/nerf/graph.py @@ -14,13 +14,24 @@ class GraphInputs: points: float +@dataclass +class Node: + """Node datastructure for graph composition.""" + + name: str + children: set + + 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 +40,39 @@ 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) -> set: + """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=set()) + 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=set([curr_module])) + processed_modules[input_module] = parent_module + else: + processed_modules[input_module].children.add(curr_module) + if input_module == "x": + roots.add(curr_module) + return roots + + 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 +80,10 @@ class Graph(nn.Module): Returns: list: ordering of the module names that should be executed """ - return modules_config.keys() + + # call construct graph + # function to do topological sort to get ordering + return self.modules_config.keys() def forward(self, x): """_summary_ @@ -49,4 +93,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..00530a164 100644 --- a/tests/nerf/test_graph.py +++ b/tests/nerf/test_graph.py @@ -4,12 +4,37 @@ 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,104 @@ 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 target_name, target_children in targets: + assert curr_pointer.name == target_name + assert len(curr_pointer.children) == target_children + if len(curr_pointer.children) > 0: + # for now traversing smaller parts of graph.. need stabler checks + curr_pointer = list(curr_pointer.children)[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}} + print(test_graph.modules_config["mlp_0"]["inputs"][1]) + 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 = [("encoder_0", 3), ("mlp_0", 2), ("mlp_3", 0)] + check_consistency(roots[0], targets) \ No newline at end of file