tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 18:44:21 +08:00
Code

rudimentary code; passing logic for ordering and input size calculation; need to check design decisions

Browse Source
This commit is contained in:
evonneng 2022-04-12 03:55:31 +00:00
parent f89bb8f7be
commit 69ea3ae45e
6 changed files with 72 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
configs/modules.json
View File

@ -4,15 +4,13 @@
"inputs": [
"x"
],
"meta_data": {
"output_dim": 8
}
"meta_data": {}
},
"mlp_0": {
"class": "MLP",
"inputs": [
"encoder_0", "encoder_0"
"encoder_0",
"encoder_0"
],
"meta_data": {
"num_layers": 8,
@ -26,10 +24,14 @@
},
"RGB": {
"name": "rgb_0",
"inputs": ["mlp_0"]
"inputs": [
"mlp_0"
]
},
"Density": {
"name": "rgb_1",
"inputs": ["mlp_0"]
"inputs": [
"mlp_0"
]
}
}

21
mattport/nerf/graph.py
View File

@ -39,6 +39,7 @@ class Graph(nn.Module):
self.roots = self.construct_graph()
self.module_order = self.get_module_order()
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"])
@ -48,26 +49,29 @@ class Graph(nn.Module):
# calculate input dimensions based on module dependencies
for root in self.roots:
self.set_in_dim(root)
self.get_in_dim(root)
def set_in_dim(self, curr_node: "Node") -> None:
# instantiate network
for _, module in self.modules.items():
module.build_model()
def get_in_dim(self, curr_node: "Node") -> None:
"""Dynamically calculates and sets the input dimensions of the modules based on dependency graph
Args:
curr_node (Node): pointer to current node in process
"""
curr_node.visited_in_dim = True
if len(curr_node.parents) > 0:
in_dim = 0
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()
self.modules[curr_node.name].set_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_order:
self.set_in_dim(child_node)
if not child_node.visited_in_dim:
self.get_in_dim(child_node)
def construct_graph(self) -> dict:
"""Constructs a dependency graph given the module configuration
@ -93,9 +97,10 @@ class Graph(nn.Module):
processed_modules[input_module] = parent_module
else:
processed_modules[input_module].children[module_name] = curr_module
curr_module.parents[input_module] = parent_module
if input_module == "x":
roots.add(curr_module)
else:
curr_module.parents[input_module] = parent_module
return roots
def topological_sort(self, curr_node: "Node", ordering_stack: list) -> None:
@ -125,7 +130,7 @@ class Graph(nn.Module):
for root in roots:
if not root.visited_order:
self.topological_sort(root, ordering_stack)
print(ordering_stack)
return ordering_stack[::-1]
def forward(self, x):

1
mattport/nerf/modules/__init__.py
View File

@ -1,3 +1,4 @@
"""init modules"""
from .encoding import Encoding
from .encoding import ScalingAndOffset
from .mlp import MLP

7
mattport/nerf/modules/encoding.py
View File

@ -19,6 +19,13 @@ class Encoding(nn.Module):
raise ValueError("Input dimension should be greater than zero")
self.in_dim = in_dim
def build_model(self) -> None:
"""TODO: create base module class from which this extends"""
def set_in_dim(self, in_dim: int) -> None:
"""TODO: create base module class from which this extends"""
self.in_dim = in_dim
def get_out_dim(self) -> int:
"""Calculates output dimension of encoding.

35
mattport/nerf/modules/mlp.py
View File

@ -19,7 +19,7 @@ class MLP(nn.Module):
activation: Optional[nn.Module] = None,
out_activation: Optional[nn.Module] = None,
) -> None:
"""Initialize mulilayer perceptron.
"""Initialize parameters of mulilayer perceptron.
Args:
in_dim (int): Input layer dimension
@ -31,20 +31,37 @@ class MLP(nn.Module):
"""
super().__init__()
self.in_dim = in_dim
self.out_dim = out_dim
self.num_layers = num_layers
self.layer_width = layer_width
self.activation = activation
self.out_activation = out_activation
self.net = None
def build_model(self) -> None:
"""Initialize mulilayer perceptron."""
layers = []
for i in range(num_layers):
for i in range(self.num_layers):
if i == 0:
layers.append(nn.Linear(in_dim, layer_width))
layers.append(nn.Linear(self.in_dim, self.layer_width))
else:
layers.append(nn.Linear(layer_width, layer_width))
if activation is not None:
layers.append(activation)
layers.append(nn.Linear(layer_width, out_dim))
if out_activation is not None:
layers.append(out_activation)
layers.append(nn.Linear(self.layer_width, self.layer_width))
if self.activation is not None:
layers.append(self.activation)
layers.append(nn.Linear(self.layer_width, self.out_dim))
if self.out_activation is not None:
layers.append(self.out_activation)
self.net = nn.Sequential(*layers)
def set_in_dim(self, in_dim: int) -> None:
"""TODO: create base module class from which this extends"""
self.in_dim = in_dim
def get_out_dim(self) -> int:
"""TODO: create base module class from which this extends"""
return self.out_dim
def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType[..., "out_dim"]:
"""Process input with a multilayer perceptron.

28
tests/nerf/test_graph.py
View File

@ -16,7 +16,7 @@ def create_basic_graph():
"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}},
"encoder_0": {"class_name": "ScalingAndOffset", "inputs": ["x"], "meta_data": {"in_dim": 16}},
"mlp_1": {
"class_name": "MLP",
"inputs": ["mlp_0"],
@ -35,7 +35,7 @@ def create_repeat_graph():
dict: module definitions for constructing a graph
"""
repeat_graph = {
"encoder_0": {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}},
"encoder_0": {"class_name": "ScalingAndOffset", "inputs": ["x"], "meta_data": {"in_dim": 16}},
"mlp_0": {
"class_name": "MLP",
"inputs": ["encoder_0", "encoder_0"],
@ -61,7 +61,7 @@ def create_complex_graph():
dict: module definitions for constructing a graph
"""
complex_graph = {
"encoder_0": {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}},
"encoder_0": {"class_name": "ScalingAndOffset", "inputs": ["x"], "meta_data": {"in_dim": 16}},
"mlp_0": {
"class_name": "MLP",
"inputs": ["encoder_0"],
@ -101,7 +101,7 @@ 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["encoder_0"]).__name__ == "ScalingAndOffset"
assert type(test_graph.modules["mlp_0"]).__name__ == "MLP"
assert type(test_graph.modules["mlp_1"]).__name__ == "MLP"
@ -109,7 +109,6 @@ def test_graph_init():
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
@ -134,7 +133,11 @@ def test_construct_graph_repeat():
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}}
test_graph.modules_config["encoder_1"] = {
"class_name": "ScalingAndOffset",
"inputs": ["x"],
"meta_data": {"out_dim": 8},
}
roots = test_graph.construct_graph()
assert len(roots) == 2
@ -162,28 +165,29 @@ 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()
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.construct_graph()
order = test_repeated_graph.get_module_order()
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": "Encoding",
"class_name": "ScalingAndOffset",
"inputs": ["x"],
"meta_data": {"out_dim": 8},
"meta_data": {"in_dim": 16},
}
order = test_repeated_graph.get_module_order()
test_repeated_graph = Graph(test_repeated_graph.modules_config)
order = test_repeated_graph.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()
order = test_complex_graph.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")