tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

Merge pull request #9 from ethanweber/modules_structuring

Browse Source 
done structuring graph + ordering module calls; [todo]: figure out calling logic using graph structure in forward()
This commit is contained in:
evonneng 2022-04-09 14:17:12 -07:00 committed by GitHub
commit 00934eed75
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 221 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

73
mattport/nerf/graph.py
View File

@ -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

162
tests/nerf/test_graph.py
View File

@ -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
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')