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 b5d2e1243..55a92e827 100644 --- a/mattport/nerf/graph.py +++ b/mattport/nerf/graph.py @@ -3,7 +3,7 @@ The Graph module contains all trainable parameters. """ import importlib from dataclasses import dataclass -from typing import Optional +from typing import Any, Dict, List, Optional, Set from torch import nn @@ -20,8 +20,10 @@ class Node: """Node datastructure for graph composition.""" name: str - children: dict - visited: Optional[bool] = False + children: Dict[str, "Node"] + parents: Dict[str, "Node"] + visited_order: Optional[bool] = False + visited_in_dim: Optional[bool] = False def __hash__(self): return hash(self.name) @@ -30,64 +32,91 @@ class Node: class Graph(nn.Module): """_summary_""" - def __init__(self, modules_config: dict) -> None: + def __init__(self, modules_config: Dict[str, Dict[str, Any]]) -> None: super().__init__() self.modules_config = modules_config - # calculate input dimensions based on module dependencies - self.modules = {} - 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"]) - # generate dependency ordering for module calls + # create graph and get ordering + self.roots = self.construct_graph() self.module_order = self.get_module_order() - def construct_graph(self) -> dict: + 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"]) + + # calculate input dimensions based on module dependencies + for root in self.roots: + self.get_in_dim(root) + + # instantiate torch.nn members of network + for _, module in self.modules.items(): + module.build_nn_modules() + + 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(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: + self.get_in_dim(child_node) + + def construct_graph(self) -> Set[Node]: """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 + set: all root nodes of 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={}) + curr_module = Node(name=module_name, children={}, parents={}) 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}) + 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 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: + def topological_sort(self, curr_node: Node, ordering_stack: List[str]) -> 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 + curr_node.visited_order = True for child_node in curr_node.children.values(): - if not child_node.visited: + if not child_node.visited_order: self.topological_sort(child_node, ordering_stack) ordering_stack.append(curr_node.name) - def get_module_order(self) -> list: + def get_module_order(self) -> List[str]: """Generates a graph and determines order of module operations using topological sorting Args: @@ -96,12 +125,11 @@ class Graph(nn.Module): Returns: list: ordering of the module names that should be executed """ - roots = self.construct_graph() + roots = self.roots ordering_stack = [] for root in roots: - if not root.visited: + if not root.visited_order: self.topological_sort(root, 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..2a96e1cc9 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 Encoding, ScalingAndOffset from .mlp import MLP +from .module import Module diff --git a/mattport/nerf/modules/encoding.py b/mattport/nerf/modules/encoding.py index 2e82ec92a..da418f406 100644 --- a/mattport/nerf/modules/encoding.py +++ b/mattport/nerf/modules/encoding.py @@ -2,11 +2,11 @@ Encoding functions """ import torch -from torch import nn from torchtyping import TensorType +from mattport.nerf.modules.module import Module -class Encoding(nn.Module): +class Encoding(Module): """Encode an input tensor. Intended to be subclassed""" def __init__(self, in_dim: int) -> None: @@ -19,14 +19,6 @@ class Encoding(nn.Module): raise ValueError("Input dimension should be greater than zero") self.in_dim = in_dim - def get_out_dim(self) -> int: - """Calculates output dimension of encoding. - - Returns: - int: output dimension - """ - 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 f278e88b9..6710742a3 100644 --- a/mattport/nerf/modules/mlp.py +++ b/mattport/nerf/modules/mlp.py @@ -5,9 +5,10 @@ from typing import Optional from torch import nn from torchtyping import TensorType +from mattport.nerf.modules.module import Module -class MLP(nn.Module): +class MLP(Module): """Multilayer perceptron""" def __init__( @@ -19,7 +20,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,17 +32,26 @@ 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) diff --git a/mattport/nerf/modules/module.py b/mattport/nerf/modules/module.py new file mode 100644 index 000000000..57dc6b5e7 --- /dev/null +++ b/mattport/nerf/modules/module.py @@ -0,0 +1,62 @@ +""" +The module baseclass. +""" +from abc import abstractmethod + +from torch import nn +from torchtyping import TensorType + + +class Module(nn.Module): + """_summary_ + + Args: + nn (_type_): _description_ + """ + + def __init__(self) -> None: + """Default initialization of module""" + super().__init__() + self.in_dim = 0 + self.out_dim = 0 + + def build_nn_modules(self) -> None: + """Function instantiates any torch.nn members within the module. + If none exist, do nothing.""" + + def set_in_dim(self, in_dim: int) -> None: + """Sets input dimension of encoding + + Args: + in_dim (int): input dimension + """ + if in_dim <= 0: + raise ValueError("Input dimension should be greater than zero") + self.in_dim = in_dim + + def get_out_dim(self) -> int: + """Calculates output dimension of encoding. + + Returns: + int: output dimension + """ + if not hasattr(self, "out_dim"): + raise ValueError("Output dimension has not been set") + if self.out_dim <= 0: + raise ValueError("Output dimension should be greater than zero") + return self.out_dim + + @abstractmethod + def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + """_summary_ + + Args: + x (TensorType[..., "input_dim"]): Input tensor to be encoded + + Raises: + NotImplementedError: _description_ + + Returns: + TensorType[..., "output_dim"]: A encoded input tensor + """ + raise NotImplementedError 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")