diff --git a/configs/modules.json b/configs/modules.json index 76c59f9d0..598925c7d 100644 --- a/configs/modules.json +++ b/configs/modules.json @@ -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" + ] } } \ No newline at end of file diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py index a59bdbd04..885da60f3 100644 --- a/mattport/nerf/graph.py +++ b/mattport/nerf/graph.py @@ -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): diff --git a/mattport/nerf/modules/__init__.py b/mattport/nerf/modules/__init__.py index f76543920..1c39985bc 100644 --- a/mattport/nerf/modules/__init__.py +++ b/mattport/nerf/modules/__init__.py @@ -1,3 +1,4 @@ """init modules""" from .encoding import Encoding +from .encoding import ScalingAndOffset from .mlp import MLP diff --git a/mattport/nerf/modules/encoding.py b/mattport/nerf/modules/encoding.py index 2e82ec92a..f06fdbdab 100644 --- a/mattport/nerf/modules/encoding.py +++ b/mattport/nerf/modules/encoding.py @@ -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. diff --git a/mattport/nerf/modules/mlp.py b/mattport/nerf/modules/mlp.py index f278e88b9..68612290b 100644 --- a/mattport/nerf/modules/mlp.py +++ b/mattport/nerf/modules/mlp.py @@ -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. diff --git a/tests/nerf/test_graph.py b/tests/nerf/test_graph.py index e73de2873..e14a89e18 100644 --- a/tests/nerf/test_graph.py +++ b/tests/nerf/test_graph.py @@ -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")