changed config loading to hydra for graph generation + tests

configs/modules.json

@@ -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"
-        ]
-    }
-}

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

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

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

tests/nerf/configs/test_repeat_graph.yaml

@@ -0,0 +1,8 @@
+defaults:
+  - test_basic_graph
+  - _self_
+  
+network:
+  mlp_0:
+    inputs: ['encoder_0', 'encoder_0']
+    

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")