From 5e24cbf6feb55bcffd89e6fd6f18f760d0acf1ba Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Fri, 15 Apr 2022 14:27:43 -0700 Subject: [PATCH 1/2] 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 From 8e852c577161e2caf6962bd1e53b406c0f0b1525 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Fri, 15 Apr 2022 14:38:26 -0700 Subject: [PATCH 2/2] Rename field heads --- .../{render_heads.py => field_outputs.py} | 12 ++++++------ mattport/nerf/renderers.py | 12 ++++++------ ...{test_render_heads.py => test_field_outputs.py} | 14 +++++++------- 3 files changed, 19 insertions(+), 19 deletions(-) rename mattport/nerf/field_modules/{render_heads.py => field_outputs.py} (90%) rename tests/nerf/field_modules/{test_render_heads.py => test_field_outputs.py} (72%) diff --git a/mattport/nerf/field_modules/render_heads.py b/mattport/nerf/field_modules/field_outputs.py similarity index 90% rename from mattport/nerf/field_modules/render_heads.py rename to mattport/nerf/field_modules/field_outputs.py index 4504cad72..262b53a00 100644 --- a/mattport/nerf/field_modules/render_heads.py +++ b/mattport/nerf/field_modules/field_outputs.py @@ -9,8 +9,8 @@ from torchtyping import TensorType from mattport.nerf.field_modules.base import FieldModule -class RenderHead(FieldModule): - """Base Render head""" +class FieldOutput(FieldModule): + """Base field output""" def __init__(self, in_dim: int, out_dim: int, activation: Optional[nn.Module] = None) -> None: """ @@ -47,15 +47,15 @@ class RenderHead(FieldModule): return self.net(in_tensor) -class DensityHead(RenderHead): - """Density head""" +class DensityFieldOutput(FieldOutput): + """Density output""" def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Softplus()) -> None: super().__init__(in_dim, 1, activation) -class RGBHead(RenderHead): - """RGB head""" +class RGBFieldOutput(FieldOutput): + """RGB output""" def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Sigmoid()) -> None: super().__init__(in_dim, 3, activation) diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index bb8b7aed6..d319ea74e 100644 --- a/mattport/nerf/renderers.py +++ b/mattport/nerf/renderers.py @@ -7,13 +7,13 @@ import torch from torch import nn from torchtyping import TensorType -from mattport.nerf.field_modules import render_heads as rh +from mattport.nerf.field_modules import field_outputs class Renderer(nn.Module): """Base Renderer. Intended to be subclassed""" - def required_render_heads(self) -> Set[rh.RenderHead]: + def required_render_heads(self) -> Set[field_outputs.FieldOutput]: """ Returns: Set[RenderHead]: RenderHeads required for this renderer. @@ -45,8 +45,8 @@ class RGB(Renderer): self.rgb_name = rgb_name self.density_name = density_name - def required_render_heads(self) -> Set[rh.RenderHead]: - return set(rh.RGBHead, rh.DensityHead) + def required_render_heads(self) -> Set[field_outputs.FieldOutput]: + return set(field_outputs.RGBFieldOutput, field_outputs.DensityFieldOutput) def forward( self, @@ -83,8 +83,8 @@ class Depth(Renderer): Renderer (_type_): _description_ """ - def required_render_heads(self) -> Set[rh.RenderHead]: - return set(rh.DensityHead) + def required_render_heads(self) -> Set[field_outputs.FieldOutput]: + return set(field_outputs.DensityFieldOutput) def forward( self, diff --git a/tests/nerf/field_modules/test_render_heads.py b/tests/nerf/field_modules/test_field_outputs.py similarity index 72% rename from tests/nerf/field_modules/test_render_heads.py rename to tests/nerf/field_modules/test_field_outputs.py index d05caeb7e..cec54dd90 100644 --- a/tests/nerf/field_modules/test_render_heads.py +++ b/tests/nerf/field_modules/test_field_outputs.py @@ -5,15 +5,15 @@ import pytest import torch from torch import nn -from mattport.nerf.field_modules import render_heads +from mattport.nerf.field_modules import field_outputs -def test_render_head(): +def test_field_output(): """Test rgb render head""" in_dim = 6 out_dim = 4 activation = nn.ReLU() - render_head = render_heads.RenderHead(in_dim=in_dim, out_dim=out_dim, activation=activation) + render_head = field_outputs.FieldOutput(in_dim=in_dim, out_dim=out_dim, activation=activation) assert render_head.get_out_dim() == out_dim x = torch.ones((9, in_dim)) @@ -27,9 +27,9 @@ def test_render_head(): def test_density_head(): - """Test rgb render head""" + """Test rgb output""" in_dim = 6 - density_head = render_heads.DensityHead(in_dim) + density_head = field_outputs.DensityFieldOutput(in_dim) assert density_head.get_out_dim() == 1 x = torch.ones((9, in_dim)) @@ -42,9 +42,9 @@ def test_density_head(): def test_rgb_head(): - """Test rgb render head""" + """Test rgb output""" in_dim = 6 - rgb_head = render_heads.RGBHead(in_dim) + rgb_head = field_outputs.RGBFieldOutput(in_dim) assert rgb_head.get_out_dim() == 3 x = torch.ones((9, in_dim))