From 39346ccf407263a815e24fb55db46b1da34611d7 Mon Sep 17 00:00:00 2001 From: evonneng Date: Tue, 12 Apr 2022 19:36:55 +0000 Subject: [PATCH] added base class and methods; done logic for building graph --- mattport/nerf/graph.py | 23 ++++++----- mattport/nerf/modules/__init__.py | 4 +- mattport/nerf/modules/encoding.py | 19 +--------- mattport/nerf/modules/mlp.py | 11 +----- mattport/nerf/modules/module.py | 63 +++++++++++++++++++++++++++++++ 5 files changed, 80 insertions(+), 40 deletions(-) create mode 100644 mattport/nerf/modules/module.py diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py index 885da60f3..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,8 @@ class Node: """Node datastructure for graph composition.""" name: str - children: dict - parents: dict + children: Dict[str, "Node"] + parents: Dict[str, "Node"] visited_order: Optional[bool] = False visited_in_dim: Optional[bool] = False @@ -32,7 +32,7 @@ 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 # create graph and get ordering @@ -51,11 +51,11 @@ class Graph(nn.Module): for root in self.roots: self.get_in_dim(root) - # instantiate network + # instantiate torch.nn members of network for _, module in self.modules.items(): - module.build_model() + module.build_nn_modules() - def get_in_dim(self, curr_node: "Node") -> None: + def get_in_dim(self, curr_node: Node) -> None: """Dynamically calculates and sets the input dimensions of the modules based on dependency graph Args: @@ -73,14 +73,14 @@ class Graph(nn.Module): if not child_node.visited_in_dim: self.get_in_dim(child_node) - def construct_graph(self) -> dict: + 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() @@ -103,7 +103,7 @@ class Graph(nn.Module): 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: @@ -116,7 +116,7 @@ class Graph(nn.Module): 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: @@ -130,7 +130,6 @@ 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 1c39985bc..2a96e1cc9 100644 --- a/mattport/nerf/modules/__init__.py +++ b/mattport/nerf/modules/__init__.py @@ -1,4 +1,4 @@ """init modules""" -from .encoding import Encoding -from .encoding import ScalingAndOffset +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 f06fdbdab..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,21 +19,6 @@ 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. - - 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 68612290b..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__( @@ -54,14 +55,6 @@ class MLP(nn.Module): 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/mattport/nerf/modules/module.py b/mattport/nerf/modules/module.py new file mode 100644 index 000000000..f7f07b279 --- /dev/null +++ b/mattport/nerf/modules/module.py @@ -0,0 +1,63 @@ +""" +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