Merge pull request #6 from ethanweber/ethan/initial_trainer_with_config

Skeleton for trainer with config

.gitignore

@@ -158,3 +158,10 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# Experiments and outputs
+runs/
+outputs/
+
+# Data
+data/

.vscode/launch.json

@@ -0,0 +1,16 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python: run_train_nerf",
+            "type": "python",
+            "request": "launch",
+            "program": "scripts/run_train_nerf.py",
+            "console": "integratedTerminal",
+            "args": []
+        }
+    ]
+}

.vscode/settings.json

@@ -13,6 +13,11 @@
     "python.linting.pylintEnabled": true,
     "python.linting.flake8Enabled": false,
     "python.linting.enabled": true,
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true,
     "python.linting.pylintArgs": [
         "--generated-members=numpy.* ,torch.* ,cv2.* , cv.*",
         "--max-line-length=120"

README.md

@@ -15,10 +15,20 @@ pip install -r environment/requirements.txt
 # Install helper library
 cd external/goat
 python setup.py develop
+# then go back, i.e., `cd /path/to/mattport`
 
 # For using with Jupyter
 python -m ipykernel install --user --name mattport --display-name "mattport"
 
 # Install mattport as a library
 python setup.py develop
+
+# Running the test cases
+pytest tests
+```
+
+# Training a model
+
+```
+python scripts/run_train_nerf.py
 ```

configs/config.yaml

@@ -0,0 +1,47 @@
+# 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:
+  data_directory: /path/to/data
+  dataset_type: blender # type of dataset
+  scene: lego
+  num_workers: 4
+
+# camera options
+camera:
+  model: pinhole
+
+# optimizer options
+optimizer:
+  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

environment/requirements.txt

@@ -1,5 +1,8 @@
+hydra-core
 pylint==2.13.4
+pytest
 torch==1.11.0
+torchtyping
 torchvision==0.12.0
 torchaudio==0.11.0
 black

mattport/nerf/dataset/base.py

@@ -0,0 +1,68 @@
+"""
+The dataset baseclass.
+"""
+
+from abc import abstractmethod
+import torch
+
+from mattport.nerf.cameras import Camera
+
+
+class Dataset(torch.utils.data.Dataset):
+    """_summary_
+
+    Args:
+        torch (_type_): _description_
+    """
+
+    def __init__(
+        self,
+        data_directory: str,
+        dataset_type: str,
+        scene: str
+    ):
+        """_summary_
+
+        Args:
+            data_directory (str): _description_
+            type (str): _description_
+            scene (str): _description_
+        """
+        super().__init__()
+        self.data_directory = data_directory
+        self.dataset_type = dataset_type
+        self.scene = scene
+
+
+    def __len__(self):
+        pass
+
+    @abstractmethod
+    def get_image(self, idx):
+        """_summary_
+
+        Args:
+            idx (_type_): _description_
+
+        Raises:
+            NotImplementedError: _description_
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_camera(self, idx) -> Camera:
+        """_summary_
+
+        Args:
+            idx (_type_): _description_
+
+        Raises:
+            NotImplementedError: _description_
+
+        Returns:
+            Camera: _description_
+        """
+        raise NotImplementedError
+
+    def __getitem__(self, idx):
+        raise NotImplementedError

mattport/nerf/dataset/build.py

@@ -0,0 +1,26 @@
+"""
+Builds the dataset
+"""
+import importlib
+import logging
+
+from mattport.nerf.dataset.base import Dataset
+
+
+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_
+    """
+
+    # TODO(ethan): better handle this with configs
+    module = importlib.import_module("mattport.nerf.dataset.base")
+    dataset = module.Dataset(
+        data_directory=config.data_directory,
+        dataset_type=config.dataset_type,
+        scene=config.scene,
+    )
+    logging.info(dataset)
+    return dataset

mattport/nerf/graph.py

@@ -0,0 +1,28 @@
+"""
+The Graph module contains all trainable parameters.
+"""
+from dataclasses import dataclass
+from torch import nn
+
+
+@dataclass
+class GraphInputs:
+    """Datastucture to encode the inputs to the graph."""
+
+    points: float
+
+
+class Graph(nn.Module):
+    """_summary_
+
+    Args:
+        nn (_type_): _description_
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        raise NotImplementedError
+
+    def forward(self):
+        """_summary_"""
+        raise NotImplementedError

mattport/nerf/optimizer.py

@@ -0,0 +1,26 @@
+"""
+Optimizer class.
+"""
+
+from torch import optim
+
+
+class Optimizer:
+    """_summary_"""
+
+    def __init__(self, params: dict, lr: float):
+        """
+        Args:
+            params (dict): _description_
+            lr (float): Learning rate.
+        """
+        self.optimizer = optim.Adam(params)
+        self.lr = lr
+
+    def test(self):
+        """_summary_
+
+        Raises:
+            NotImplementedError: _description_
+        """
+        raise NotImplementedError

mattport/nerf/trainer.py

@@ -0,0 +1,61 @@
+"""
+Code to train model.
+"""
+from mattport.nerf.dataset.build import build_dataset
+from mattport.nerf.optimizer import Optimizer
+from mattport.nerf.graph import Graph
+
+
+class Trainer(object):
+    """Training class"""
+
+    def __init__(self, config: dict):
+        self.config = config
+        self.train_dataset = None
+        self.test_dataset = None
+        self.graph = None
+        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()
+
+    def setup_optimizer(self):
+        """_summary_
+        """
+        self.optimizer = Optimizer(params=self.graph.parameters(), **self.config.optimizer)
+
+    def load_checkpoint(self):
+        """_summary_
+
+        Raises:
+            NotImplementedError: _description_
+        """
+        raise NotImplementedError
+
+    def save_checkpoint(self):
+        """_summary_
+
+        Raises:
+            NotImplementedError: _description_
+        """
+        raise NotImplementedError
+
+    def train(self):
+        """_summary_"""
+        raise NotImplementedError
+
+    def train_epoch(self):
+        """_summary_"""
+        raise NotImplementedError
+
+    def train_iteration(self):
+        """_summary_"""
+        raise NotImplementedError

mattport/viewer/README.md

@@ -6,4 +6,5 @@ The backend viewer code (python, websockets, ...).
 
 ```
 https://github.com/lp-research/openvslam/tree/main/viewer
+https://github.com/rdeits/meshcat
 ```

scripts/run_train_nerf.py

@@ -1,5 +1,24 @@
-import os
-import mattport.nerf as nerf
+"""
+run_train_nerf.py
+"""
+import logging
+import hydra
+from omegaconf import DictConfig
+
+from mattport.nerf.trainer import Trainer
+
+
+logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG)
+
+
+@hydra.main(config_path="../configs", config_name="config.yaml")
+def main(config: DictConfig):
+    """Main function."""
+    trainer = Trainer(config)
+    trainer.setup_dataset()
+    trainer.setup_graph()
+    trainer.setup_optimizer()
+
 
 if __name__ == "__main__":
-    print("hello mattport")
+    main()

tests/nerf/test_encodings.py

@@ -2,4 +2,8 @@ import pytest
 
 
 def test_positional_encoding():
-    raise NotImplementedError()
+    assert True
+
+
+def test_lower_case():
+    assert "MATTPORT".lower() == "mattport"