brainstorming trainer code but broken

.gitignore

@@ -161,4 +161,7 @@ cython_debug/
 
 # Experiments and outputs
 runs/
-outputs/
+outputs/
+
+# Data
+data/

configs/config.yaml

@@ -1,6 +1,50 @@
-defaults:
-  - experiment: lego_bulldozer
-  - dataset: blender
-batch_size: 10
+# defaults:
+#   - experiment: lego_bulldozer
+#   - dataset: blender
+
+# OPTIONS FOR DIFFERENT NERF VARIANTS
+# ----
+# sampler options
+sampler:
+  asdf: null
+  num_coarse_samples: 64
+  num_importance_samples: 128
+  perturb: 1.0
+  noise_std: 1.0
+# encoding options
+encodings:
+  asdf: null
+# network options
+network:
+  asdf: null
+# renderer options
+renderer:
+  asdf: null
+# loss options
+loss:
+  asdf: null
+  type: mse
+# ----
+
+# dataset options
+dataset:
+  dataset: blender # type of dataset
+  scene: lego
+  num_workers: 4
+
+# camera options
+camera:
+  model: pinhole
+
+# optimizer options
 optimizer:
-  - lr: 1e-4
+  lr: 1e-4
+  lr_scheduler: null
+  momentum: null
+  weight_decay: null
+
+# trainer options
+trainer:
+  batch_size: 2 # rays per batch
+  chunk_size: 1024 # chunks per batch
+  device: cpu

mattport/nerf/dataset/base.py

@@ -0,0 +1,31 @@
+from abc import abstractmethod
+import torch
+
+
+class Dataset(torch.utils.data.Dataset):
+    """_summary_
+
+    Args:
+        torch (_type_): _description_
+    """
+
+    def __init__(
+        self,
+        data_directory: str,
+    ) -> None:
+        super().__init__()
+        pass
+
+    def __len__(self):
+        pass
+
+    @abstractmethod
+    def get_image(self, idx):
+        pass
+
+    @abstractmethod
+    def get_pose(self, idx):
+        pass
+
+    def __getitem__(self, idx):
+        raise NotImplementedError

mattport/nerf/dataset/build.py

@@ -0,0 +1,17 @@
+from argon2 import DEFAULT_RANDOM_SALT_LENGTH
+from mattport.nerf.dataset.base import Dataset
+import importlib
+import logging
+
+def build_dataset(config: dict) -> Dataset:
+    """Returns the dataset according to the specified config.
+    # TODO(ethan): change the type from dict to be a dataclass
+
+    Args:
+        config (dict): _description_
+    """
+
+    module = importlib.import_module("mattport.nerf.dataset.base")
+    dataset = module.Dataset(**config)
+    logging.info(dataset)
+    return dataset

mattport/nerf/graph.py

@@ -0,0 +1,15 @@
+from torch import nn
+
+
+class Graph(nn.Module):
+    """_summary_
+
+    Args:
+        nn (_type_): _description_
+    """
+
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.config = config
+
+    

mattport/nerf/optimizer.py

@@ -1,10 +1,18 @@
-import torch.optim as optim
+from torch import optim
 
 
 class Optimizer:
     """_summary_"""
 
-    def __init__(self, params: dict) -> None:
+    def __init__(
+        self, 
+        params: dict,
+        lr: float) -> None:
+        """
+        Args:
+            params (dict): _description_
+            lr (float): Learning rate.
+        """
         self.optimizer = optim.Adam(params)
 
     def test(self):

mattport/nerf/trainer.py

@@ -1,11 +1,33 @@
+from abc import abstractmethod
 import logging
 import torch
 
+from mattport.nerf.optimizer import Optimizer
+from mattport.nerf.graph import Graph
+from mattport.nerf.dataset.build import build_dataset
+
 
 class Trainer(object):
-    def __init__(self, config: None):
+    def __init__(self, config: dict):
         self.config = config
-        self.optimizer = None
+
+    def setup_dataset(self):
+        """_summary_
+        """
+        self.train_dataset = build_dataset(self.config.dataset)
+        self.test_dataset = build_dataset(self.config.dataset)
+
+    def setup_graph(self):
+        """_summary_
+        """
+        self.graph = Graph(
+            encodings_config=self.config.encodings,
+            renderer_config=self.config.renderer,
+            loss_config=self.config.loss
+        )
+
+    def setup_optimizer(self):
+        self.optimizer = Optimizer(params=self.graph.parameters(), **self.config.optimizer)
 
     def load_checkpoint(self):
         """_summary_
@@ -13,7 +35,7 @@ class Trainer(object):
         Raises:
             NotImplementedError: _description_
         """
-        raise NotImplementedError()
+        raise NotImplementedError
 
     def save_checkpoint(self):
         """_summary_
@@ -21,17 +43,16 @@ class Trainer(object):
         Raises:
             NotImplementedError: _description_
         """
-        raise NotImplementedError()
-
-    def load_dataset(self):
-        """_summary_
-
-        Raises:
-            NotImplementedError: _description_
-        """
-        raise NotImplementedError()
+        raise NotImplementedError
 
     def train(self):
         """_summary_"""
-        # TODO(ethan): replace print statements with something more proper
-        logging.info("train function")
+        pass
+
+    def train_epoch(self):
+        """_summary_"""
+        pass
+
+    def train_iteration(self):
+        """_summary_"""
+        pass

scripts/run_train_nerf.py

@@ -9,17 +9,13 @@ logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=loggin
 
 
 @hydra.main(config_path="../configs", config_name="config.yaml")
-def main(cfg: DictConfig):
-    print(cfg)
-
-    print("hello mattport")
-
-    # TODO(ethan): add some sort of config script
-
-    config = {}
+def main(config: DictConfig):
     trainer = Trainer(config)
     print("trainer!")
-    trainer.train()
+
+    trainer.setup_dataset()
+    # trainer.setup_graph()
+    # trainer.setup_optimizer()
 
 
 if __name__ == "__main__":