diff --git a/configs/modules.json b/configs/modules.json new file mode 100644 index 000000000..377de3a0d --- /dev/null +++ b/configs/modules.json @@ -0,0 +1,34 @@ +{ + "PositionalEncoder": { + "name": "encoder_0", + "inputs": [ + "x" + ], + "meta_data": { + "output_dim": 8 + } + }, + "MLP": { + "name": "mlp_0", + "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"] + } +} \ No newline at end of file diff --git a/mattport/nerf/dataset/build.py b/mattport/nerf/dataset/build.py index 6966e87e6..6e0d9c4b3 100644 --- a/mattport/nerf/dataset/build.py +++ b/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: diff --git a/mattport/nerf/dataset/base.py b/mattport/nerf/dataset/dataset.py similarity index 100% rename from mattport/nerf/dataset/base.py rename to mattport/nerf/dataset/dataset.py diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py index 01d1a6db1..ee0cf3efc 100644 --- a/mattport/nerf/graph.py +++ b/mattport/nerf/graph.py @@ -3,7 +3,7 @@ The Graph module contains all trainable parameters. """ from dataclasses import dataclass from torch import nn - +from mattport.nerf.modules.positional_encoder import PositionalEncoder @dataclass class GraphInputs: @@ -19,10 +19,31 @@ class Graph(nn.Module): nn (_type_): _description_ """ - def __init__(self) -> None: + 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.encoder_0 = Encoder() + self.mlp_0 = MLP() + + raise NotImplementedError - def forward(self): + def forward(self, x): """_summary_""" - raise NotImplementedError + # raise NotImplementedError + + for name in queue: # "encoder_0", + # assert name.output is not None + x = self.modules[name](x) + # self.mlp_0 \ No newline at end of file diff --git a/mattport/nerf/encodings.py b/mattport/nerf/modules/encoding.py similarity index 92% rename from mattport/nerf/encodings.py rename to mattport/nerf/modules/encoding.py index e25c3a50e..98daf6180 100644 --- a/mattport/nerf/encodings.py +++ b/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, name: str, inputs: list, meta_data: dict) -> None: + super().__init__() + raise NotImplementedError + def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: """Encodes an input tensor. diff --git a/mattport/nerf/models.py b/mattport/nerf/modules/mlp.py similarity index 94% rename from mattport/nerf/models.py rename to mattport/nerf/modules/mlp.py index f58ecc717..94ac56e0e 100644 --- a/mattport/nerf/models.py +++ b/mattport/nerf/modules/mlp.py @@ -1,10 +1,4 @@ -""" -Collection of Models -""" -from typing import Optional - from torch import nn -from torchtyping import TensorType class MLP(nn.Module): @@ -30,7 +24,6 @@ class MLP(nn.Module): out_activation (Optional[nn.Module], optional): output activation function. Defaults to None. """ super().__init__() - layers = [] for i in range(num_layers): if i == 0: diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 23c479f97..13e58d663 100644 --- a/mattport/nerf/trainer.py +++ b/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):