brainstorming with matt on vanilla nerf implementation

mattport/nerf/field/nerf.py

@@ -0,0 +1,52 @@
+"""
+Multi Layer Perceptron
+"""
+from typing import Optional
+
+from torch import nn
+from torchtyping import TensorType
+from mattport.nerf.field_modules.encoding import NeRFEncoding
+from mattport.nerf.field_modules.field_outputs import DensityFieldOutput, RGBFieldOutput
+from mattport.nerf.field_modules.mlp import MLP
+
+
+class NeRFField():
+    """Multilayer perceptron"""
+
+    def __init__(self) -> None:
+        self.build_nn_modules()
+
+    def build_nn_modules(self) -> None:
+        """Initialize mulilayer perceptron."""
+        self.encoding_xyz = NeRFEncoding(in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0)
+        self.encoding_dir = NeRFEncoding(in_dim=3, num_frequencies=6, min_freq_exp=0.0, max_freq_exp=4.0)
+        self.mlp_base = MLP(
+            in_dim=self.encoding_xyz.get_out_dim(), out_dim=64, num_layers=8, layer_width=64, activation=nn.ReLU()
+        )
+        self.mlp_rgb = MLP(
+            in_dim=self.base_mlp.get_out_dim() + self.encoding_dir.get_out_dim(),
+            out_dim=64,
+            num_layers=2,
+            layer_width=64,
+            activation=nn.ReLU(),
+        )
+        self.field_output_rgb = RGBFieldOutput(in_dim=self.mlp_rgb.get_out_dim())
+        self.field_output_density = DensityFieldOutput(in_dim=self.mlp_base.get_out_dim())
+
+    def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType[..., "out_dim"]:
+        """Process input with a multilayer perceptron.
+
+        Args:
+            in_tensor (TensorType[..., "in_dim]): Network input
+
+        Returns:
+            TensorType[..., "out_dim"]: Network output
+        """
+
+        encoded_xyz = self.encoding_xyz(xyz)
+        encoded_dir = self.encoding_dir(rays.direction)
+        base_mlp_out = self.mlp_base(encoded_xyz)
+        rgb_mlp_out = self.mlp_rgb(torch.cat([encoded_dir, base_mlp_out], dim=-1))
+        field_rgb_output = self.field_output_rgb(rgb_mlp_out)
+        field_density_out = self.field_output_density(base_mlp_out)
+        return rgb_coarse

mattport/nerf/field_modules/base.py

@@ -19,6 +19,7 @@ class FieldModule(nn.Module):
     def build_nn_modules(self) -> None:
         """Function instantiates any torch.nn members within the module.
         If none exist, do nothing."""
+        raise ValueError("This function shouldn't be called.")
 
     def set_in_dim(self, in_dim: int) -> None:
         """Sets input dimension of encoding

mattport/nerf/field_modules/field_outputs.py

@@ -23,8 +23,7 @@ class FieldOutput(FieldModule):
         self.in_dim = in_dim
         self.out_dim = out_dim
         self.activation = activation
-
-        self.net = None
+        self.net = self.build_nn_modules()
 
     def build_nn_modules(self) -> None:
         layers = [nn.Linear(self.in_dim, self.out_dim)]

mattport/nerf/field_modules/mlp.py

@@ -37,7 +37,7 @@ class MLP(FieldModule):
         self.layer_width = layer_width
         self.activation = activation
         self.out_activation = out_activation
-        self.net = None
+        self.net = self.build_nn_modules()
 
     def build_nn_modules(self) -> None:
         """Initialize mulilayer perceptron."""

mattport/nerf/graph/base.py

@@ -4,6 +4,7 @@ The Graph module contains all trainable parameters.
 import importlib
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Set
+from black import Priority
 from omegaconf import DictConfig
 
 from torch import nn
@@ -146,4 +147,10 @@ class Graph(nn.ModuleDict):
         #     x = self.modules[name](x)
 
         ## can index into previous results using dictionary
-        raise NotImplementedError
+        # raise NotImplementedError
+        for inputs, node_name, out_name in priority_queue:
+            self.outputs[out_name] = self.modules[node_name](inputs)
+
+            # possibly delete things from output if no longer needed
+            for old_name in :
+                del self.outputs[old_name]

mattport/nerf/graph/nerf.py

@@ -0,0 +1,79 @@
+"""
+Implementation of vanilla nerf.
+"""
+
+
+import torch
+from omegaconf import DictConfig
+from torch import TensorType, nn
+
+from mattport.nerf.field_modules.encoding import NeRFEncoding
+from mattport.nerf.field_modules.field_outputs import DensityFieldOutput, RGBFieldOutput
+from mattport.nerf.field_modules.mlp import MLP
+from mattport.nerf.field.nerf import NeRFField
+from mattport.nerf.field_modules.ray_generator import RayGenerator
+from mattport.nerf.graph.base import Graph
+from mattport.nerf.loss import MSELoss
+from mattport.nerf.renderers import RGBRenderer
+from mattport.nerf.sampler import UniformSampler
+from mattport.structures.cameras import Rays
+from mattport.nerf.dataset.utils import DatasetInputs
+
+
+class NeRFGraph(Graph):
+    def __init__(self, modules_config: DictConfig, dataset_inputs: DatasetInputs) -> None:
+        super().__init__(modules_config)
+
+        self.populate_modules(dataset_inputs=dataset_inputs)
+        # TODO(ethan): build all modules here instead of in the code below
+
+    def populate_modules(self, dataset_inputs: DatasetInputs):
+
+        # ray generator
+        self.ray_generator = RayGenerator(dataset_inputs.intrinsics, dataset_inputs.camera_to_world)
+
+        # samplers
+        self.sampler_uniform = UniformSampler(near_plane=0.1, far_plane=4.0, num_samples=64)
+        self.sampler_pdf = UniformSampler(near_plane=0.1, far_plane=4.0, num_samples=64)
+
+        # field
+        self.coarse = NeRFField()
+        self.fine = NeRFField()
+
+        # renderers
+        self.renderer_rgb = RGBRenderer()
+
+        # losses
+        self.rgb_loss = MSELoss()
+
+    def forward(self, batch_indices: TensorType["num_rays", 3]):
+
+        ray_bundle = self.ray_generator.forward(batch_indices)
+        rays = Rays(origin=ray_bundle.origins, direction=ray_bundle.directions)  # TODO(ethan): consolidate this
+        ray_samples = self.sampler_uniform(rays)
+
+        xyz = ray_samples.locations
+        deltas = ray_samples.deltas
+
+        encoded_xyz = self.encoding_xyz(xyz)
+        encoded_dir = self.encoding_dir(rays.direction)
+        base_mlp_out = self.mlp_base(encoded_xyz)
+        rgb_mlp_out = self.mlp_rgb(torch.cat([encoded_dir, base_mlp_out], dim=-1))
+        field_rgb_output = self.field_output_rgb(rgb_mlp_out)
+        field_density_out = self.field_output_density(base_mlp_out)
+        rgb_coarse = self.renderer_rgb(rgb=field_rgb_output, density=field_density_out, deltas=deltas)
+
+        xyz_pdf = self.sampler_pdf(rays)  # TODO(ethan): implement pdf
+        encoded_xyz_pdf = self.encoding_xyz(xyz_pdf)
+        base_mlp_out = self.mlp_base(encoded_xyz_pdf)
+        rgb_mlp_out = self.mlp_rgb(torch.cat([encoded_dir, base_mlp_out], dim=-1))
+        field_rgb_output = self.field_output_rgb(rgb_mlp_out)
+        field_density_out = self.field_output_density(base_mlp_out)
+        rgb_fine = self.renderer_rgb(rgb=field_rgb_output, density=field_density_out, deltas=deltas)
+
+        return {"rgb_coarse": rgb_coarse, "rgb_fine": rgb_fine}
+
+    def get_loss(self, batch, render_outputs):
+        rgb_loss_coarse = self.rgb_loss(batch.rgb, render_outputs["rgb_coarse"])
+        rgb_loss_fine = self.rgb_loss(batch.rgb, render_outputs["rgb_fine"])
+        return {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine}

mattport/nerf/renderers.py

@@ -37,20 +37,19 @@ class Renderer(nn.Module):
         raise NotImplementedError
 
 
-class RGB(Renderer):
+class RGBRenderer(Renderer):
     """Standard volumetic rendering."""
 
-    def __init__(self, rgb_name, density_name) -> None:
+    def __init__(self) -> None:
         super().__init__()
-        self.rgb_name = rgb_name
-        self.density_name = density_name
 
     def required_field_outputs(self) -> Set[field_outputs.FieldOutput]:
         return set(field_outputs.RGBFieldOutput, field_outputs.DensityFieldOutput)
 
     def forward(
         self,
-        field_output_dict: Dict[str, TensorType[..., "num_samples", -1]],
+        rgb: TensorType[..., "num_samples", 3],
+        density: TensorType[..., "num_samples", 1],
         deltas: TensorType[..., "num_samples"],
     ) -> TensorType[..., "out_dim"]:
         """Composite samples along ray and render color image
@@ -62,7 +61,7 @@ class RGB(Renderer):
         Returns:
             TensorType[..., "out_dim"]: Composited RGB ray
         """
-        delta_density = deltas * field_output_dict[self.density_name][..., 0]
+        delta_density = deltas * density[..., 0]
         alphas = 1 - torch.exp(-delta_density)
 
         transmittance = torch.cumsum(delta_density[..., :-1], dim=-1)
@@ -71,12 +70,12 @@ class RGB(Renderer):
 
         weights = alphas * transmittance  # [..., "num_samples"]
 
-        rgb = torch.sum(weights[..., None] * field_output_dict[self.rgb_name], dim=-2)
+        rgb = torch.sum(weights[..., None] * rgb, dim=-2)
 
         return rgb
 
 
-class Depth(Renderer):
+class DepthRenderer(Renderer):
     """_summary_
 
     Args:

mattport/nerf/trainer.py

@@ -1,11 +1,15 @@
 """
 Code to train model.
 """
-from omegaconf import DictConfig
 import torch.distributed as dist
+from omegaconf import DictConfig
 from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader
+from tqdm import tqdm
 
-from mattport.nerf.graph import Graph
+from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch
+from mattport.nerf.dataset.utils import get_dataset_inputs
+from mattport.nerf.graph.base import Graph
 from mattport.nerf.optimizer import Optimizer
 
 
@@ -23,8 +27,20 @@ class Trainer:
 
     def setup_dataset(self):
         """_summary_"""
-        self.train_dataset = None
-        self.test_dataset = None
+        dataset_inputs = get_dataset_inputs(**self.config.dataset)
+        self.train_dataset = ImageDataset(
+            image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor
+        )
+        self.train_dataloader = DataLoader(
+            self.train_dataset,
+            batch_size=self.config.dataloader.num_images_to_sample_from,
+            collate_fn=lambda batch: collate_batch(
+                batch, self.config.dataloader.num_rays_per_batch, keep_full_image=True
+            ),
+            num_workers=self.config.dataloader.num_workers,
+            shuffle=True,
+        )
+        # TODO(ethan): implement test data
 
     def setup_graph(self):
         """_summary_"""
@@ -60,10 +76,17 @@ class Trainer:
 
     def train_epoch(self):
         """_summary_"""
-        raise NotImplementedError
+        for _ in range(len(self.train_dataset)):
+            batch = next(iter(self.train_dataloader))
+            self.train_iteration(batch)
 
-    def train_iteration(self):
+    def train_iteration(self, batch):
         """_summary_"""
         # TODO(): save checkpoints, and do logging only on rank 0 device only
         # if gpu == 0:
-        raise NotImplementedError
+        output_dict = self.graph(batch.indices)
+        
+        loss = self.graph.get_loss(batch, )
+
+        loss.backward()
+        # update model