diff --git a/configs/modules.json b/configs/modules.json index 377de3a0d..76c59f9d0 100644 --- a/configs/modules.json +++ b/configs/modules.json @@ -1,6 +1,6 @@ { - "PositionalEncoder": { - "name": "encoder_0", + "encoder_0": { + "class": "PositionalEncoder", "inputs": [ "x" ], @@ -8,8 +8,9 @@ "output_dim": 8 } }, - "MLP": { - "name": "mlp_0", + + "mlp_0": { + "class": "MLP", "inputs": [ "encoder_0", "encoder_0" ], diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py index ee0cf3efc..843cde7c2 100644 --- a/mattport/nerf/graph.py +++ b/mattport/nerf/graph.py @@ -1,9 +1,11 @@ """ The Graph module contains all trainable parameters. """ +import importlib from dataclasses import dataclass + from torch import nn -from mattport.nerf.modules.positional_encoder import PositionalEncoder + @dataclass class GraphInputs: @@ -21,29 +23,35 @@ class Graph(nn.Module): def __init__(self, modules_config: dict) -> None: super().__init__() - - - # # for key in ... - # temp = graph_config.positional_encoder - # PositionalEncoder(name=temp.name, name=temp.inputs, ) - # self.modules = {} - # self.modules["name"] = MLP() - - 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": + in_dim = 0 + for inputs in module_dict["inputs"]: + 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"]) + def get_module_order(self, modules_config: list): + """Generates a graph to determine the order in which the graph should execute - self.encoder_0 = Encoder() - self.mlp_0 = MLP() + Args: + modules_config (list): _description_ - - raise NotImplementedError + Returns: + _type_: _description_ + """ + # raise NotImplementedError + return modules_config def forward(self, x): - """_summary_""" + """_summary_ + + Args: + x (_type_): _description_ + """ # raise NotImplementedError - - for name in queue: # "encoder_0", - # assert name.output is not None + for name in self.module_order: x = self.modules[name](x) - # self.mlp_0 \ No newline at end of file diff --git a/mattport/nerf/modules/__init__.py b/mattport/nerf/modules/__init__.py new file mode 100644 index 000000000..d59de93cb --- /dev/null +++ b/mattport/nerf/modules/__init__.py @@ -0,0 +1,2 @@ +from .encoding import Encoding +from .mlp import MLP \ No newline at end of file diff --git a/mattport/nerf/modules/encoding.py b/mattport/nerf/modules/encoding.py index 98daf6180..821c663c5 100644 --- a/mattport/nerf/modules/encoding.py +++ b/mattport/nerf/modules/encoding.py @@ -8,9 +8,9 @@ from torchtyping import TensorType class Encoding(nn.Module): """Encode an input tensor. Intended to be subclassed""" - def __init__(self, name: str, inputs: list, meta_data: dict) -> None: + def __init__(self, out_dim: int) -> None: super().__init__() - raise NotImplementedError + # raise NotImplementedError def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: """Encodes an input tensor. diff --git a/mattport/nerf/modules/mlp.py b/mattport/nerf/modules/mlp.py index 94ac56e0e..6c47264da 100644 --- a/mattport/nerf/modules/mlp.py +++ b/mattport/nerf/modules/mlp.py @@ -1,4 +1,7 @@ +from typing import Optional + from torch import nn +from torchtyping import TensorType class MLP(nn.Module): @@ -10,7 +13,7 @@ class MLP(nn.Module): out_dim: int, num_layers: int, layer_width: int, - activation: Optional[nn.Module] = nn.ReLU, + activation: Optional[nn.Module] = None, out_activation: Optional[nn.Module] = None, ) -> None: """Initialize mulilayer perceptron. @@ -24,6 +27,7 @@ class MLP(nn.Module): out_activation (Optional[nn.Module], optional): output activation function. Defaults to None. """ super().__init__() + self.in_dim = in_dim layers = [] for i in range(num_layers): if i == 0: diff --git a/tests/nerf/test_graph.py b/tests/nerf/test_graph.py new file mode 100644 index 000000000..5dbf9e468 --- /dev/null +++ b/tests/nerf/test_graph.py @@ -0,0 +1,43 @@ +import pytest + +from mattport.nerf.graph import Graph + + +def create_basic_graph(): + """Create basic test for creating a graph + + Returns: + dict: module definitions for constructing a graph + """ + basic_graph = { + "encoder_0": {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}}, + "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 basic_graph + + +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' + + +def test_input_dimension(): + """test module input dimension calculation""" + 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 + \ No newline at end of file