From 5e24cbf6feb55bcffd89e6fd6f18f760d0acf1ba Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Fri, 15 Apr 2022 14:27:43 -0700 Subject: [PATCH] Restructure into field modules --- mattport/nerf/{modules => field_modules}/__init__.py | 4 ++-- .../nerf/{modules/module.py => field_modules/base.py} | 10 +++------- mattport/nerf/{modules => field_modules}/encoding.py | 4 ++-- mattport/nerf/field_modules/field_losses.py | 0 mattport/nerf/{modules => field_modules}/mlp.py | 4 ++-- .../nerf/{modules => field_modules}/ray_generator.py | 0 .../nerf/{modules => field_modules}/render_heads.py | 4 ++-- mattport/nerf/graph.py | 2 +- mattport/nerf/renderers.py | 2 +- .../nerf/{modules => field_modules}/test_encodings.py | 2 +- .../{modules => field_modules}/test_render_heads.py | 2 +- tests/nerf/test_graph.py | 8 +++++--- tests/nerf/test_ray_generator.py | 2 +- 13 files changed, 21 insertions(+), 23 deletions(-) rename mattport/nerf/{modules => field_modules}/__init__.py (56%) rename mattport/nerf/{modules/module.py => field_modules/base.py} (91%) rename mattport/nerf/{modules => field_modules}/encoding.py (98%) create mode 100644 mattport/nerf/field_modules/field_losses.py rename mattport/nerf/{modules => field_modules}/mlp.py (96%) rename mattport/nerf/{modules => field_modules}/ray_generator.py (100%) rename mattport/nerf/{modules => field_modules}/render_heads.py (95%) rename tests/nerf/{modules => field_modules}/test_encodings.py (97%) rename tests/nerf/{modules => field_modules}/test_render_heads.py (95%) diff --git a/mattport/nerf/modules/__init__.py b/mattport/nerf/field_modules/__init__.py similarity index 56% rename from mattport/nerf/modules/__init__.py rename to mattport/nerf/field_modules/__init__.py index 2a96e1cc9..80c8bcdb2 100644 --- a/mattport/nerf/modules/__init__.py +++ b/mattport/nerf/field_modules/__init__.py @@ -1,4 +1,4 @@ -"""init modules""" +"""init field modules""" from .encoding import Encoding, ScalingAndOffset from .mlp import MLP -from .module import Module +from .base import FieldModule diff --git a/mattport/nerf/modules/module.py b/mattport/nerf/field_modules/base.py similarity index 91% rename from mattport/nerf/modules/module.py rename to mattport/nerf/field_modules/base.py index 57dc6b5e7..b3c0a8e52 100644 --- a/mattport/nerf/modules/module.py +++ b/mattport/nerf/field_modules/base.py @@ -1,5 +1,5 @@ """ -The module baseclass. +The field module baseclass. """ from abc import abstractmethod @@ -7,12 +7,8 @@ from torch import nn from torchtyping import TensorType -class Module(nn.Module): - """_summary_ - - Args: - nn (_type_): _description_ - """ +class FieldModule(nn.Module): + """Field modules that can be combined to store and compute the fields.""" def __init__(self) -> None: """Default initialization of module""" diff --git a/mattport/nerf/modules/encoding.py b/mattport/nerf/field_modules/encoding.py similarity index 98% rename from mattport/nerf/modules/encoding.py rename to mattport/nerf/field_modules/encoding.py index da418f406..751370c6b 100644 --- a/mattport/nerf/modules/encoding.py +++ b/mattport/nerf/field_modules/encoding.py @@ -3,10 +3,10 @@ Encoding functions """ import torch from torchtyping import TensorType -from mattport.nerf.modules.module import Module +from mattport.nerf.field_modules.base import FieldModule -class Encoding(Module): +class Encoding(FieldModule): """Encode an input tensor. Intended to be subclassed""" def __init__(self, in_dim: int) -> None: diff --git a/mattport/nerf/field_modules/field_losses.py b/mattport/nerf/field_modules/field_losses.py new file mode 100644 index 000000000..e69de29bb diff --git a/mattport/nerf/modules/mlp.py b/mattport/nerf/field_modules/mlp.py similarity index 96% rename from mattport/nerf/modules/mlp.py rename to mattport/nerf/field_modules/mlp.py index 6710742a3..008920d97 100644 --- a/mattport/nerf/modules/mlp.py +++ b/mattport/nerf/field_modules/mlp.py @@ -5,10 +5,10 @@ from typing import Optional from torch import nn from torchtyping import TensorType -from mattport.nerf.modules.module import Module +from mattport.nerf.field_modules.base import FieldModule -class MLP(Module): +class MLP(FieldModule): """Multilayer perceptron""" def __init__( diff --git a/mattport/nerf/modules/ray_generator.py b/mattport/nerf/field_modules/ray_generator.py similarity index 100% rename from mattport/nerf/modules/ray_generator.py rename to mattport/nerf/field_modules/ray_generator.py diff --git a/mattport/nerf/modules/render_heads.py b/mattport/nerf/field_modules/render_heads.py similarity index 95% rename from mattport/nerf/modules/render_heads.py rename to mattport/nerf/field_modules/render_heads.py index 1c2b5f114..4504cad72 100644 --- a/mattport/nerf/modules/render_heads.py +++ b/mattport/nerf/field_modules/render_heads.py @@ -6,10 +6,10 @@ from typing import Optional from torch import nn from torchtyping import TensorType -from mattport.nerf.modules.module import Module +from mattport.nerf.field_modules.base import FieldModule -class RenderHead(Module): +class RenderHead(FieldModule): """Base Render head""" def __init__(self, in_dim: int, out_dim: int, activation: Optional[nn.Module] = None) -> None: diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py index cd6108080..5dd87440d 100644 --- a/mattport/nerf/graph.py +++ b/mattport/nerf/graph.py @@ -43,7 +43,7 @@ class Graph(nn.Module): self.modules = {} # initialize graph with known input dimensions; set default in_dim to 0 for module_name, module_dict in modules_config.items(): - module = getattr(importlib.import_module("mattport.nerf.modules"), module_dict.class_name) + module = getattr(importlib.import_module("mattport.nerf.field_modules"), module_dict.class_name) if not module_dict.meta_data.in_dim: module_dict.meta_data.in_dim = 0 self.modules[module_name] = module(**module_dict.meta_data) diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index fa9f83ee3..bb8b7aed6 100644 --- a/mattport/nerf/renderers.py +++ b/mattport/nerf/renderers.py @@ -7,7 +7,7 @@ import torch from torch import nn from torchtyping import TensorType -from mattport.nerf.modules import render_heads as rh +from mattport.nerf.field_modules import render_heads as rh class Renderer(nn.Module): diff --git a/tests/nerf/modules/test_encodings.py b/tests/nerf/field_modules/test_encodings.py similarity index 97% rename from tests/nerf/modules/test_encodings.py rename to tests/nerf/field_modules/test_encodings.py index 47bd74668..6798ead49 100644 --- a/tests/nerf/modules/test_encodings.py +++ b/tests/nerf/field_modules/test_encodings.py @@ -3,7 +3,7 @@ Encoding Tests """ import pytest import torch -from mattport.nerf.modules import encoding +from mattport.nerf.field_modules import encoding def test_scaling_and_offset(): diff --git a/tests/nerf/modules/test_render_heads.py b/tests/nerf/field_modules/test_render_heads.py similarity index 95% rename from tests/nerf/modules/test_render_heads.py rename to tests/nerf/field_modules/test_render_heads.py index b22c48499..d05caeb7e 100644 --- a/tests/nerf/modules/test_render_heads.py +++ b/tests/nerf/field_modules/test_render_heads.py @@ -5,7 +5,7 @@ import pytest import torch from torch import nn -from mattport.nerf.modules import render_heads +from mattport.nerf.field_modules import render_heads def test_render_head(): diff --git a/tests/nerf/test_graph.py b/tests/nerf/test_graph.py index 73d8f0a25..e45665341 100644 --- a/tests/nerf/test_graph.py +++ b/tests/nerf/test_graph.py @@ -1,5 +1,7 @@ +""" +Tests of graph construction +""" import numpy as np -import pytest from hydra import compose, initialize from omegaconf import open_dict @@ -28,7 +30,7 @@ def instantiate_graph(config_name: str) -> Graph: def test_graph_init(): - """test modules dictionary is instantiated properly""" + """test field modules dictionary is instantiated properly""" test_graph, _ = instantiate_graph("test_basic_graph") assert test_graph.modules is not None assert type(test_graph.modules["encoder_0"]).__name__ == "ScalingAndOffset" @@ -37,7 +39,7 @@ def test_graph_init(): def test_input_dimension(): - """test module input dimension calculation""" + """test field module input dimension calculation""" test_graph, _ = instantiate_graph("test_repeat_graph") assert test_graph.modules["mlp_0"].in_dim == 16 assert test_graph.modules["mlp_1"].in_dim == 1 diff --git a/tests/nerf/test_ray_generator.py b/tests/nerf/test_ray_generator.py index 22ba267cc..e9d3bb36b 100644 --- a/tests/nerf/test_ray_generator.py +++ b/tests/nerf/test_ray_generator.py @@ -4,7 +4,7 @@ Test ray generation. import torch from mattport.nerf.dataset.blender_dataset import load_blender_data -from mattport.nerf.modules.ray_generator import RayGenerator +from mattport.nerf.field_modules.ray_generator import RayGenerator from mattport.structures.rays import RayBundle from mattport.utils.io import get_absolute_path