From a9c0d5b97f51670e176e6ebd159764a028e1e1a8 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Wed, 6 Apr 2022 16:46:47 +0000 Subject: [PATCH 1/3] hydra config and starting trainer class --- .gitignore | 4 +++ .vscode/launch.json | 16 +++++++++++ .vscode/settings.json | 5 ++++ README.md | 10 +++++++ configs/config.yaml | 6 +++++ configs/dataset/blender.yaml | 0 configs/dataset/dtu.yaml | 0 configs/dataset/shapenet.yaml | 0 configs/experiment/lego_bulldozer.yaml | 2 ++ configs/optimizer/adam.yaml | 0 configs/optimizer/sgd.yaml | 0 environment/requirements.txt | 3 +++ mattport/nerf/optimizer.py | 13 +++++++++ mattport/nerf/trainer.py | 37 ++++++++++++++++++++++++++ mattport/viewer/README.md | 1 + scripts/run_train_nerf.py | 25 +++++++++++++++-- tests/nerf/test_encodings.py | 6 ++++- 17 files changed, 125 insertions(+), 3 deletions(-) create mode 100644 .vscode/launch.json create mode 100644 configs/config.yaml create mode 100644 configs/dataset/blender.yaml create mode 100644 configs/dataset/dtu.yaml create mode 100644 configs/dataset/shapenet.yaml create mode 100644 configs/experiment/lego_bulldozer.yaml create mode 100644 configs/optimizer/adam.yaml create mode 100644 configs/optimizer/sgd.yaml create mode 100644 mattport/nerf/optimizer.py diff --git a/.gitignore b/.gitignore index 68bc17f9f..efc8ee8f3 100644 --- a/.gitignore +++ b/.gitignore @@ -158,3 +158,7 @@ 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/ \ No newline at end of file diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 000000000..13aa61c4c --- /dev/null +++ b/.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": [] + } + ] +} \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json index 8c50b4b35..3f6d1df2c 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -13,4 +13,9 @@ "python.linting.pylintEnabled": true, "python.linting.flake8Enabled": false, "python.linting.enabled": true, + "python.testing.pytestArgs": [ + "tests" + ], + "python.testing.unittestEnabled": false, + "python.testing.pytestEnabled": true, } \ No newline at end of file diff --git a/README.md b/README.md index f0be6b824..092eea60a 100644 --- a/README.md +++ b/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 ``` \ No newline at end of file diff --git a/configs/config.yaml b/configs/config.yaml new file mode 100644 index 000000000..864783428 --- /dev/null +++ b/configs/config.yaml @@ -0,0 +1,6 @@ +defaults: + - experiment: lego_bulldozer + - dataset: blender +batch_size: 10 +optimizer: + - lr: 1e-4 diff --git a/configs/dataset/blender.yaml b/configs/dataset/blender.yaml new file mode 100644 index 000000000..e69de29bb diff --git a/configs/dataset/dtu.yaml b/configs/dataset/dtu.yaml new file mode 100644 index 000000000..e69de29bb diff --git a/configs/dataset/shapenet.yaml b/configs/dataset/shapenet.yaml new file mode 100644 index 000000000..e69de29bb diff --git a/configs/experiment/lego_bulldozer.yaml b/configs/experiment/lego_bulldozer.yaml new file mode 100644 index 000000000..02e9326c7 --- /dev/null +++ b/configs/experiment/lego_bulldozer.yaml @@ -0,0 +1,2 @@ +name: lego_bulldozer +feature_size: 5 \ No newline at end of file diff --git a/configs/optimizer/adam.yaml b/configs/optimizer/adam.yaml new file mode 100644 index 000000000..e69de29bb diff --git a/configs/optimizer/sgd.yaml b/configs/optimizer/sgd.yaml new file mode 100644 index 000000000..e69de29bb diff --git a/environment/requirements.txt b/environment/requirements.txt index c1b62abdf..ef64b7b0c 100644 --- a/environment/requirements.txt +++ b/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 diff --git a/mattport/nerf/optimizer.py b/mattport/nerf/optimizer.py new file mode 100644 index 000000000..f2bd07230 --- /dev/null +++ b/mattport/nerf/optimizer.py @@ -0,0 +1,13 @@ +import torch.optim as optim + + +class Optimizer: + """_summary_""" + + def __init__(self, params: dict) -> None: + self.optimizer = optim.Adam(params) + + def test(self): + raise NotImplementedError + + \ No newline at end of file diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index e69de29bb..3de4b3918 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -0,0 +1,37 @@ +import logging +import torch + + +class Trainer(object): + def __init__(self, config: None): + self.config = config + self.optimizer = None + + def load_checkpoint(self): + """_summary_ + + Raises: + NotImplementedError: _description_ + """ + raise NotImplementedError() + + def save_checkpoint(self): + """_summary_ + + Raises: + NotImplementedError: _description_ + """ + raise NotImplementedError() + + def load_dataset(self): + """_summary_ + + Raises: + NotImplementedError: _description_ + """ + raise NotImplementedError() + + def train(self): + """_summary_""" + # TODO(ethan): replace print statements with something more proper + logging.info("train function") diff --git a/mattport/viewer/README.md b/mattport/viewer/README.md index aca76ca0b..3ba37fcd0 100644 --- a/mattport/viewer/README.md +++ b/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 ``` \ No newline at end of file diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py index aaa3ffc7b..1bd96d3c0 100644 --- a/scripts/run_train_nerf.py +++ b/scripts/run_train_nerf.py @@ -1,5 +1,26 @@ import os -import mattport.nerf as nerf +from typing import Dict +from mattport.nerf.trainer import Trainer +import logging +import hydra +from omegaconf import DictConfig + +logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG) + + +@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 = {} + trainer = Trainer(config) + print("trainer!") + trainer.train() + if __name__ == "__main__": - print("hello mattport") + main() diff --git a/tests/nerf/test_encodings.py b/tests/nerf/test_encodings.py index 49ad82215..99f26096c 100644 --- a/tests/nerf/test_encodings.py +++ b/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" From a7167f890e97d5d37f05c415ca4428579c7ced25 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Wed, 6 Apr 2022 22:45:53 +0000 Subject: [PATCH 2/3] 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__": From 5691d00d427c2f3e42688e387a8f228fbb4f0aaf Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Thu, 7 Apr 2022 20:22:48 +0000 Subject: [PATCH 3/3] fixed most linting errors --- .vscode/settings.json | 1 + configs/config.yaml | 7 ++-- configs/dataset/blender.yaml | 0 configs/dataset/dtu.yaml | 0 configs/dataset/shapenet.yaml | 0 configs/experiment/lego_bulldozer.yaml | 2 -- configs/optimizer/adam.yaml | 0 configs/optimizer/sgd.yaml | 0 mattport/nerf/dataset/base.py | 47 +++++++++++++++++++++++--- mattport/nerf/dataset/build.py | 17 +++++++--- mattport/nerf/graph.py | 19 +++++++++-- mattport/nerf/optimizer.py | 17 ++++++---- mattport/nerf/trainer.py | 29 +++++++++------- scripts/run_train_nerf.py | 16 +++++---- 14 files changed, 110 insertions(+), 45 deletions(-) delete mode 100644 configs/dataset/blender.yaml delete mode 100644 configs/dataset/dtu.yaml delete mode 100644 configs/dataset/shapenet.yaml delete mode 100644 configs/experiment/lego_bulldozer.yaml delete mode 100644 configs/optimizer/adam.yaml delete mode 100644 configs/optimizer/sgd.yaml diff --git a/.vscode/settings.json b/.vscode/settings.json index 3f6d1df2c..f476de807 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -11,6 +11,7 @@ "--line-length=120" ], "python.linting.pylintEnabled": true, + "python.linting.pylintArgs": [], "python.linting.flake8Enabled": false, "python.linting.enabled": true, "python.testing.pytestArgs": [ diff --git a/configs/config.yaml b/configs/config.yaml index 2175cf59b..f95fc0ce2 100644 --- a/configs/config.yaml +++ b/configs/config.yaml @@ -1,7 +1,3 @@ -# defaults: -# - experiment: lego_bulldozer -# - dataset: blender - # OPTIONS FOR DIFFERENT NERF VARIANTS # ---- # sampler options @@ -28,7 +24,8 @@ loss: # dataset options dataset: - dataset: blender # type of dataset + data_directory: /path/to/data + dataset_type: blender # type of dataset scene: lego num_workers: 4 diff --git a/configs/dataset/blender.yaml b/configs/dataset/blender.yaml deleted file mode 100644 index e69de29bb..000000000 diff --git a/configs/dataset/dtu.yaml b/configs/dataset/dtu.yaml deleted file mode 100644 index e69de29bb..000000000 diff --git a/configs/dataset/shapenet.yaml b/configs/dataset/shapenet.yaml deleted file mode 100644 index e69de29bb..000000000 diff --git a/configs/experiment/lego_bulldozer.yaml b/configs/experiment/lego_bulldozer.yaml deleted file mode 100644 index 02e9326c7..000000000 --- a/configs/experiment/lego_bulldozer.yaml +++ /dev/null @@ -1,2 +0,0 @@ -name: lego_bulldozer -feature_size: 5 \ No newline at end of file diff --git a/configs/optimizer/adam.yaml b/configs/optimizer/adam.yaml deleted file mode 100644 index e69de29bb..000000000 diff --git a/configs/optimizer/sgd.yaml b/configs/optimizer/sgd.yaml deleted file mode 100644 index e69de29bb..000000000 diff --git a/mattport/nerf/dataset/base.py b/mattport/nerf/dataset/base.py index 67cbd6422..2c48232fa 100644 --- a/mattport/nerf/dataset/base.py +++ b/mattport/nerf/dataset/base.py @@ -1,6 +1,12 @@ +""" +The dataset baseclass. +""" + from abc import abstractmethod import torch +from mattport.nerf.cameras import Camera + class Dataset(torch.utils.data.Dataset): """_summary_ @@ -12,20 +18,51 @@ class Dataset(torch.utils.data.Dataset): def __init__( self, data_directory: str, - ) -> None: + dataset_type: str, + scene: str + ): + """_summary_ + + Args: + data_directory (str): _description_ + type (str): _description_ + scene (str): _description_ + """ super().__init__() - pass + self.data_directory = data_directory + self.dataset_type = dataset_type + self.scene = scene + def __len__(self): pass @abstractmethod def get_image(self, idx): - pass + """_summary_ + + Args: + idx (_type_): _description_ + + Raises: + NotImplementedError: _description_ + """ + raise NotImplementedError @abstractmethod - def get_pose(self, idx): - pass + 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 diff --git a/mattport/nerf/dataset/build.py b/mattport/nerf/dataset/build.py index e9eceb00f..6966e87e6 100644 --- a/mattport/nerf/dataset/build.py +++ b/mattport/nerf/dataset/build.py @@ -1,8 +1,12 @@ -from argon2 import DEFAULT_RANDOM_SALT_LENGTH -from mattport.nerf.dataset.base import Dataset +""" +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 @@ -11,7 +15,12 @@ def build_dataset(config: dict) -> Dataset: config (dict): _description_ """ + # TODO(ethan): better handle this with configs module = importlib.import_module("mattport.nerf.dataset.base") - dataset = module.Dataset(**config) + dataset = module.Dataset( + data_directory=config.data_directory, + dataset_type=config.dataset_type, + scene=config.scene, + ) logging.info(dataset) - return dataset \ No newline at end of file + return dataset diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py index ec52c9850..01d1a6db1 100644 --- a/mattport/nerf/graph.py +++ b/mattport/nerf/graph.py @@ -1,6 +1,17 @@ +""" +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_ @@ -8,8 +19,10 @@ class Graph(nn.Module): nn (_type_): _description_ """ - def __init__(self, config) -> None: + def __init__(self) -> None: super().__init__() - self.config = config + raise NotImplementedError - \ No newline at end of file + def forward(self): + """_summary_""" + raise NotImplementedError diff --git a/mattport/nerf/optimizer.py b/mattport/nerf/optimizer.py index 3484bada3..e6c6581cf 100644 --- a/mattport/nerf/optimizer.py +++ b/mattport/nerf/optimizer.py @@ -1,21 +1,26 @@ +""" +Optimizer class. +""" + from torch import optim class Optimizer: """_summary_""" - def __init__( - self, - params: dict, - lr: float) -> None: + 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): - raise NotImplementedError + """_summary_ - \ No newline at end of file + Raises: + NotImplementedError: _description_ + """ + raise NotImplementedError diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 4722d16a8..23c479f97 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -1,15 +1,20 @@ -from abc import abstractmethod -import logging -import torch - +""" +Code to train model. +""" +from mattport.nerf.dataset.build import build_dataset from mattport.nerf.optimizer import Optimizer from mattport.nerf.graph import Graph -from mattport.nerf.dataset.build import build_dataset 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_ @@ -20,13 +25,11 @@ class Trainer(object): def setup_graph(self): """_summary_ """ - self.graph = Graph( - encodings_config=self.config.encodings, - renderer_config=self.config.renderer, - loss_config=self.config.loss - ) + self.graph = Graph() def setup_optimizer(self): + """_summary_ + """ self.optimizer = Optimizer(params=self.graph.parameters(), **self.config.optimizer) def load_checkpoint(self): @@ -47,12 +50,12 @@ class Trainer(object): def train(self): """_summary_""" - pass + raise NotImplementedError def train_epoch(self): """_summary_""" - pass + raise NotImplementedError def train_iteration(self): """_summary_""" - pass + raise NotImplementedError diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py index 65a097aef..2ad9a5f97 100644 --- a/scripts/run_train_nerf.py +++ b/scripts/run_train_nerf.py @@ -1,21 +1,23 @@ -import os -from typing import Dict -from mattport.nerf.trainer import Trainer +""" +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) - print("trainer!") - trainer.setup_dataset() - # trainer.setup_graph() - # trainer.setup_optimizer() + trainer.setup_graph() + trainer.setup_optimizer() if __name__ == "__main__":