From a7167f890e97d5d37f05c415ca4428579c7ced25 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Wed, 6 Apr 2022 22:45:53 +0000 Subject: [PATCH] brainstorming trainer code but broken --- .gitignore | 5 ++- configs/config.yaml | 54 ++++++++++++++++++++++++++++--- mattport/nerf/dataset/__init__.py | 0 mattport/nerf/dataset/base.py | 31 ++++++++++++++++++ mattport/nerf/dataset/build.py | 17 ++++++++++ mattport/nerf/graph.py | 15 +++++++++ mattport/nerf/optimizer.py | 12 +++++-- mattport/nerf/trainer.py | 49 ++++++++++++++++++++-------- scripts/run_train_nerf.py | 14 +++----- 9 files changed, 166 insertions(+), 31 deletions(-) create mode 100644 mattport/nerf/dataset/__init__.py create mode 100644 mattport/nerf/dataset/base.py create mode 100644 mattport/nerf/dataset/build.py create mode 100644 mattport/nerf/graph.py diff --git a/.gitignore b/.gitignore index efc8ee8f3..de34d8e54 100644 --- a/.gitignore +++ b/.gitignore @@ -161,4 +161,7 @@ cython_debug/ # Experiments and outputs runs/ -outputs/ \ No newline at end of file +outputs/ + +# Data +data/ \ No newline at end of file diff --git a/configs/config.yaml b/configs/config.yaml index 864783428..2175cf59b 100644 --- a/configs/config.yaml +++ b/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 \ No newline at end of file diff --git a/mattport/nerf/dataset/__init__.py b/mattport/nerf/dataset/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/mattport/nerf/dataset/base.py b/mattport/nerf/dataset/base.py new file mode 100644 index 000000000..67cbd6422 --- /dev/null +++ b/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 diff --git a/mattport/nerf/dataset/build.py b/mattport/nerf/dataset/build.py new file mode 100644 index 000000000..e9eceb00f --- /dev/null +++ b/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 \ No newline at end of file diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py new file mode 100644 index 000000000..ec52c9850 --- /dev/null +++ b/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 + + \ No newline at end of file diff --git a/mattport/nerf/optimizer.py b/mattport/nerf/optimizer.py index f2bd07230..3484bada3 100644 --- a/mattport/nerf/optimizer.py +++ b/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): diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 3de4b3918..4722d16a8 100644 --- a/mattport/nerf/trainer.py +++ b/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 diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py index 1bd96d3c0..65a097aef 100644 --- a/scripts/run_train_nerf.py +++ b/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__":