diff --git a/configs/modules.json b/configs/modules.json deleted file mode 100644 index 598925c7d..000000000 --- a/configs/modules.json +++ /dev/null @@ -1,37 +0,0 @@ -{ - "encoder_0": { - "class": "PositionalEncoder", - "inputs": [ - "x" - ], - "meta_data": {} - }, - "mlp_0": { - "class": "MLP", - "inputs": [ - "encoder_0", - "encoder_0" - ], - "meta_data": { - "num_layers": 8, - "num_skip_connections": [ - 3, - 7 - ], - "hidden_dim": 256, - "output_dim": 8 - } - }, - "RGB": { - "name": "rgb_0", - "inputs": [ - "mlp_0" - ] - }, - "Density": { - "name": "rgb_1", - "inputs": [ - "mlp_0" - ] - } -} \ No newline at end of file diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py index 55a92e827..cd6108080 100644 --- a/mattport/nerf/graph.py +++ b/mattport/nerf/graph.py @@ -3,7 +3,8 @@ The Graph module contains all trainable parameters. """ import importlib from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Set +from typing import Dict, List, Optional, Set +from omegaconf import DictConfig from torch import nn @@ -32,7 +33,7 @@ class Node: class Graph(nn.Module): """_summary_""" - def __init__(self, modules_config: Dict[str, Dict[str, Any]]) -> None: + def __init__(self, modules_config: DictConfig) -> None: super().__init__() self.modules_config = modules_config # create graph and get ordering @@ -42,10 +43,10 @@ class Graph(nn.Module): self.modules = {} # 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 "in_dim" not in module_dict["meta_data"]: - module_dict["meta_data"]["in_dim"] = 0 - self.modules[module_name] = module(**module_dict["meta_data"]) + module = getattr(importlib.import_module("mattport.nerf.modules"), module_dict.class_name) + if not module_dict.meta_data.in_dim: + module_dict.meta_data.in_dim = 0 + self.modules[module_name] = module(**module_dict.meta_data) # calculate input dimensions based on module dependencies for root in self.roots: @@ -67,7 +68,7 @@ class Graph(nn.Module): 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(in_dim) - self.modules_config[curr_node.name]["meta_data"]["in_dim"] = 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_in_dim: @@ -90,10 +91,14 @@ class Graph(nn.Module): processed_modules[module_name] = curr_module else: curr_module = processed_modules[module_name] - inputs = module_dict["inputs"] + 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}, parents={}) + 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 diff --git a/tests/nerf/configs/__init__.py b/tests/nerf/configs/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/nerf/configs/test_basic_graph.yaml b/tests/nerf/configs/test_basic_graph.yaml new file mode 100644 index 000000000..b3e23bd79 --- /dev/null +++ b/tests/nerf/configs/test_basic_graph.yaml @@ -0,0 +1,22 @@ +network: + 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: ScalingAndOffset + inputs: ['x'] + meta_data: + in_dim: 16 + mlp_1: + class_name: MLP + inputs: ['mlp_0'] + meta_data: + in_dim: 4 + out_dim: 5 + num_layers: 6 + layer_width: 7 \ No newline at end of file diff --git a/tests/nerf/configs/test_complex_graph.yaml b/tests/nerf/configs/test_complex_graph.yaml new file mode 100644 index 000000000..276fd6bcb --- /dev/null +++ b/tests/nerf/configs/test_complex_graph.yaml @@ -0,0 +1,23 @@ +defaults: + - test_basic_graph + - _self_ + +network: + mlp_1: + inputs: ['encoder_0'] + 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 \ No newline at end of file diff --git a/tests/nerf/configs/test_repeat_graph.yaml b/tests/nerf/configs/test_repeat_graph.yaml new file mode 100644 index 000000000..3abb91e6f --- /dev/null +++ b/tests/nerf/configs/test_repeat_graph.yaml @@ -0,0 +1,8 @@ +defaults: + - test_basic_graph + - _self_ + +network: + mlp_0: + inputs: ['encoder_0', 'encoder_0'] + \ No newline at end of file diff --git a/tests/nerf/test_graph.py b/tests/nerf/test_graph.py index e14a89e18..42356ff50 100644 --- a/tests/nerf/test_graph.py +++ b/tests/nerf/test_graph.py @@ -1,92 +1,14 @@ import numpy as np import pytest -from mattport.nerf.graph import Graph +from hydra import compose, initialize +from omegaconf import open_dict + +from mattport.nerf.graph import Graph, Node + +CONFIG_DIR = "configs" -def create_basic_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": "ScalingAndOffset", "inputs": ["x"], "meta_data": {"in_dim": 16}}, - "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": "ScalingAndOffset", "inputs": ["x"], "meta_data": {"in_dim": 16}}, - "mlp_0": { - "class_name": "MLP", - "inputs": ["encoder_0", "encoder_0"], - "meta_data": {"in_dim": 0, "out_dim": 1, "num_layers": 2, "layer_width": 3}, - }, - "mlp_1": { - "class_name": "MLP", - "inputs": ["mlp_0"], - "meta_data": {"in_dim": 4, "out_dim": 5, "num_layers": 6, "layer_width": 7}, - }, - } - 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": "ScalingAndOffset", "inputs": ["x"], "meta_data": {"in_dim": 16}}, - "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): +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): @@ -97,9 +19,17 @@ def check_consistency(root: "Node", targets: list): curr_pointer = curr_pointer.children[targets[i + 1][0]] +def instantiate_graph(config_name: str) -> Graph: + """load correct graph given hydra config""" + with initialize(config_path=CONFIG_DIR): + cfg = compose(config_name=config_name) + test_graph = Graph(cfg.network) + return test_graph, cfg + + def test_graph_init(): """test modules dictionary is instantiated properly""" - test_graph = Graph(create_basic_graph()) + test_graph, _ = instantiate_graph("test_basic_graph") assert test_graph.modules is not None assert type(test_graph.modules["encoder_0"]).__name__ == "ScalingAndOffset" assert type(test_graph.modules["mlp_0"]).__name__ == "MLP" @@ -108,14 +38,14 @@ def test_graph_init(): def test_input_dimension(): """test module input dimension calculation""" - test_graph = Graph(create_repeat_graph()) + test_graph, _ = instantiate_graph("test_repeat_graph") 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()) + test_graph, _ = instantiate_graph("test_basic_graph") roots = test_graph.construct_graph() assert len(roots) == 1 roots = list(roots) @@ -125,26 +55,28 @@ def test_construct_graph_basic(): def test_construct_graph_repeat(): """test construction of dependency graph: ensure duplicate nodes are ok""" - test_graph = Graph(create_repeat_graph()) + test_graph, cfg = instantiate_graph("test_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": "ScalingAndOffset", - "inputs": ["x"], - "meta_data": {"out_dim": 8}, - } + with open_dict(cfg): + cfg.network.mlp_0.inputs[1] = "encoder_1" + cfg.network.encoder_1 = { + "class_name": "ScalingAndOffset", + "inputs": ["x"], + "meta_data": {"in_dim": 8}, + } + test_graph = Graph(cfg.network) 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()) + test_graph, _ = instantiate_graph("test_complex_graph") roots = test_graph.construct_graph() assert len(roots) == 1 roots = list(roots) @@ -163,29 +95,32 @@ def test_construct_graph_complex(): def test_ordering(): """test dependency ordering of constructed graph""" - test_basic_graph = Graph(create_basic_graph()) + test_basic_graph, _ = instantiate_graph("test_basic_graph") _ = test_basic_graph.construct_graph() order = test_basic_graph.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, cfg = instantiate_graph("test_repeat_graph") _ = test_repeated_graph.construct_graph() order = test_repeated_graph.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": "ScalingAndOffset", - "inputs": ["x"], - "meta_data": {"in_dim": 16}, - } - test_repeated_graph = Graph(test_repeated_graph.modules_config) + with open_dict(cfg): + cfg.network.mlp_0.inputs[1] = "encoder_1" + cfg.network.encoder_1 = { + "class_name": "ScalingAndOffset", + "inputs": ["x"], + "meta_data": {"in_dim": 16}, + } + test_repeated_graph = Graph(cfg.network) order = test_repeated_graph.module_order - assert order.index("mlp_0") > order.index("encoder_0") and order.index("mlp_0") > order.index("encoder_1") + 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, _ = instantiate_graph("test_complex_graph") _ = test_complex_graph.construct_graph() order = test_complex_graph.module_order assert order.index("mlp_2") > order.index("mlp_0")