Merge pull request #8 from ethanweber/modules_structuring

modules initialization logic

.github/workflows/pytest.yml

@@ -0,0 +1,31 @@
+name: Python application
+
+on:
+  push:
+    branches: [ master ]
+  pull_request:
+    branches: [ master ]
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v3
+      with:
+        python-version: "3.10"
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install flake8 pytest
+        if [ -f environment/requirements.txt ]; then pip install -r environment/requirements.txt; fi
+        python setup.py develop
+    - name: Test with pytest
+      run: |
+        pytest tests/nerf/test_graph.py

configs/modules.json

@@ -0,0 +1,35 @@
+{
+    "encoder_0": {
+        "class": "PositionalEncoder",
+        "inputs": [
+            "x"
+        ],
+        "meta_data": {
+            "output_dim": 8
+        }
+    },
+
+    "mlp_0": {
+        "class": "MLP",
+        "inputs": [
+            "encoder_0", "encoder_0"
+        ],
+        "meta_data": {
+            "num_layers": 8,
+            "num_skip_connections": [
+                3,
+                7
+            ],
+            "hidden_dim": 256,
+            "output_dim": 8
+        }
+    },
+    "RGB": {
+        "name": "rgb_0",
+        "inputs": ["mlp_0"]
+    },
+    "Density": {
+        "name": "rgb_1",
+        "inputs": ["mlp_0"]
+    }
+}

mattport/nerf/dataset/build.py

@@ -4,7 +4,7 @@ Builds the dataset
 import importlib
 import logging
 
-from mattport.nerf.dataset.base import Dataset
+from mattport.nerf.dataset.dataset import Dataset
 
 
 def build_dataset(config: dict) -> Dataset:

mattport/nerf/graph.py

@@ -1,7 +1,9 @@
 """
 The Graph module contains all trainable parameters.
 """
+import importlib
 from dataclasses import dataclass
+
 from torch import nn
 
 
@@ -13,16 +15,38 @@ class GraphInputs:
 
 
 class Graph(nn.Module):
-    """_summary_
+    """_summary_"""
 
-    Args:
-        nn (_type_): _description_
-    """
-
-    def __init__(self) -> None:
+    def __init__(self, modules_config: dict) -> None:
         super().__init__()
-        raise NotImplementedError
+        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 forward(self):
-        """_summary_"""
+    def get_module_order(self, modules_config: dict) -> list:
+        """Generates a graph to determine the order in which the modules should execute
+
+        Args:
+            modules_config (dict): module definitions that make up the network
+
+        Returns:
+            list: ordering of the module names that should be executed
+        """
+        return modules_config.keys()
+
+    def forward(self, x):
+        """_summary_
+
+        Args:
+            x (_type_): _description_
+        """
+        # for name in self.module_order:
+        #     x = self.modules[name](x)
         raise NotImplementedError

mattport/nerf/modules/__init__.py

@@ -0,0 +1,2 @@
+from .encoding import Encoding
+from .mlp import MLP

mattport/nerf/modules/encoding.py

@@ -8,6 +8,10 @@ from torchtyping import TensorType
 class Encoding(nn.Module):
     """Encode an input tensor. Intended to be subclassed"""
 
+    def __init__(self, out_dim: int) -> None:
+        super().__init__()
+        self.out_dim = out_dim
+
     def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
         """Encodes an input tensor.
 

mattport/nerf/modules/mlp.py

@@ -1,6 +1,3 @@
-"""
-Collection of Models
-"""
 from typing import Optional
 
 from torch import nn
@@ -16,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.
@@ -30,7 +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:

mattport/nerf/trainer.py

@@ -17,19 +17,16 @@ class Trainer(object):
         self.optimizer = None
 
     def setup_dataset(self):
-        """_summary_
-        """
+        """_summary_"""
         self.train_dataset = build_dataset(self.config.dataset)
         self.test_dataset = build_dataset(self.config.dataset)
 
     def setup_graph(self):
-        """_summary_
-        """
+        """_summary_"""
         self.graph = Graph()
 
     def setup_optimizer(self):
-        """_summary_
-        """
+        """_summary_"""
         self.optimizer = Optimizer(params=self.graph.parameters(), **self.config.optimizer)
 
     def load_checkpoint(self):

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
+