network pass running

mattport/nerf/field/nerf.py

@@ -1,16 +1,14 @@
 """
 Multi Layer Perceptron
 """
-from typing import Optional
 
 import torch
 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.field_heads import DensityFieldHead, FieldHeadOutputs, RGBFieldHead
 from mattport.nerf.field_modules.mlp import MLP
-from mattport.structures.rays import RayBundle, RaySamples
+from mattport.structures.rays import RaySamples
 
 
 class NeRFField(nn.Module):
@@ -34,22 +32,26 @@ class NeRFField(nn.Module):
             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, xyz: TensorType["num_rays", 3], dirs: TensorType["num_rays", 3]):
-        """
+        self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_rgb.get_out_dim())
+        self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim())
 
+    def forward(self, ray_samples: RaySamples):
+        """Evaluates the field at points along the ray
         Args:
             xyz: ()
         # TODO(ethan): change the input to be something more abstracted
         e.g., a FieldInput structure
         """
-
-        encoded_xyz = self.encoding_xyz(xyz)
-        encoded_dir = self.encoding_dir(dirs)
+        positions = ray_samples.positions
+        directions = ray_samples.directions
+        encoded_xyz = self.encoding_xyz(positions)
+        encoded_dir = self.encoding_dir(directions)
         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 field_rgb_output, field_density_out
+
+        # TODO(ethan): implement a way to combine the field outputs...
+        field_outputs = FieldHeadOutputs(rgb=field_rgb_output.rgb, density=field_density_out.density)
+        return field_outputs

mattport/nerf/field_modules/base.py

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

mattport/nerf/field_modules/field_heads.py

@@ -1,6 +1,7 @@
 """
 Collection of render heads
 """
+from dataclasses import dataclass
 from typing import Optional
 
 from torch import nn
@@ -9,7 +10,15 @@ from torchtyping import TensorType
 from mattport.nerf.field_modules.base import FieldModule
 
 
-class FieldOutput(FieldModule):
+@dataclass
+class FieldHeadOutputs:
+    """_summary_"""
+
+    rgb: TensorType["...", 3] = None
+    density: TensorType["...", 1] = None
+
+
+class FieldHead(FieldModule):
     """Base field output"""
 
     def __init__(self, in_dim: int, out_dim: int, activation: Optional[nn.Module] = None) -> None:
@@ -41,20 +50,26 @@ class FieldOutput(FieldModule):
         Returns:
             TensorType[..., "out_dim"]: Render head output
         """
+        if not self.field_quantity_name:
+            raise ValueError("field_quantity_name should be set in the child class. E.g., as 'rgb' or 'density'.")
         if not self.net:
             raise SystemError("Render head network not initialized. build_nn_modules() should be called.")
-        return self.net(in_tensor)
+        out_tensor = self.net(in_tensor)
+        field_head_outputs = FieldHeadOutputs()
+        setattr(field_head_outputs, self.field_quantity_name, out_tensor)
+        return field_head_outputs
 
-
-class DensityFieldOutput(FieldOutput):
+class DensityFieldHead(FieldHead):
     """Density output"""
 
     def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Softplus()) -> None:
         super().__init__(in_dim, 1, activation)
+        self.field_quantity_name = "density"
 
 
-class RGBFieldOutput(FieldOutput):
+class RGBFieldHead(FieldHead):
     """RGB output"""
 
     def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Sigmoid()) -> None:
         super().__init__(in_dim, 3, activation)
+        self.field_quantity_name = "rgb"

mattport/nerf/graph/base.py

@@ -1,16 +1,13 @@
 """
 The Graph module contains all trainable parameters.
 """
-from abc import abstractmethod
 import importlib
+from abc import abstractmethod
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Set
-from black import Priority
-from omegaconf import DictConfig
 
 from torch import nn
-
-from mattport.nerf.dataset.utils import DatasetInputs
+from torchtyping import TensorType
 
 
 @dataclass
@@ -68,7 +65,6 @@ class Graph(nn.ModuleDict):
     @abstractmethod
     def populate_modules(self):
         """Initializes the modules that are part of the network."""
-        pass
 
     def get_in_dim(self, curr_node: Node) -> None:
         """Dynamically calculates and sets the input dimensions of the modules based on dependency graph
@@ -152,11 +148,9 @@ class Graph(nn.ModuleDict):
         return ordering_stack[::-1]
 
     @abstractmethod
-    def forward(self):
+    def forward(self, ray_indices: TensorType["num_rays", 3]):
         """Forward function that needs to be overridden."""
-        pass
 
     @abstractmethod
-    def get_losses(self):
+    def get_losses(self, batch, graph_outputs):
         """Computes and returns the losses."""
-        pass

mattport/nerf/graph/nerf.py

@@ -3,18 +3,9 @@ Implementation of vanilla nerf.
 """
 
 
-import logging
-
-import torch
-from omegaconf import DictConfig
-from torch import nn
 from torchtyping import TensorType
 
-from mattport.nerf.dataset.utils import DatasetInputs, get_dataset_inputs
 from mattport.nerf.field.nerf import NeRFField
-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_modules.ray_generator import RayGenerator
 from mattport.nerf.graph.base import Graph
 from mattport.nerf.loss import MSELoss
@@ -23,6 +14,11 @@ from mattport.nerf.sampler import UniformSampler
 
 
 class NeRFGraph(Graph):
+    """_summary_
+
+    Args:
+        Graph (_type_): _description_
+    """
     def __init__(self, intrinsics=None, camera_to_world=None) -> None:
         super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world)
 
@@ -46,36 +42,28 @@ class NeRFGraph(Graph):
 
     def forward(self, ray_indices: TensorType["num_rays", 3]):
         """Takes in the ray indices and renders out values."""
-
-        ray_bundle = self.ray_generator.forward(ray_indices)
-        ray_samples = self.sampler_uniform(ray_bundle)
-
-        xyz = ray_samples.get_positions(ray_bundle) # (num_rays, num_samples, 3)
-        num_rays, num_samples = xyz.shape[:2]
-        xyz = xyz.view(-1, 3)
-        dirs = (
-            ray_bundle.directions.unsqueeze(1).repeat(1, num_samples, 1).view(num_rays * num_samples, 3)
-        )  # (num_rays, 3)
-        
-        field_output_rgb, field_output_density = self.field_coarse(xyz=xyz, dirs=dirs)
-        # move these back to the correct shape
-        field_output_rgb = field_output_rgb.view(num_rays, num_samples, 3)
-        field_output_density = field_output_density.view(num_rays, num_samples)
-
-        deltas = ray_samples.get_deltas()
-        rgb_coarse = self.renderer_rgb(rgb=field_output_rgb, density=field_output_density, deltas=deltas)
-        
-        print(field_output_rgb)
-        print(field_output_density)
-        outputs = {
-            "rgb_coarse": rgb_coarse
-        }
+        # get the rays:
+        ray_bundle = self.ray_generator.forward(ray_indices)  # RayBundle
+        # coarse network:
+        uniform_ray_samples = self.sampler_uniform(ray_bundle)  # RaySamples
+        coarse_field_outputs = self.field_coarse(uniform_ray_samples)  # FieldOutputs
+        coarse_renderer_outputs = self.renderer_rgb(
+            rgb=coarse_field_outputs.rgb, density=coarse_field_outputs.density, deltas=uniform_ray_samples.deltas
+        )  # RendererOutputs
+        # fine network:
+        pdf_ray_samples = self.sampler_uniform(ray_bundle)  # RaySamples
+        fine_field_outputs = self.field_fine(pdf_ray_samples)  # FieldOutputs
+        fine_renderer_outputs = self.renderer_rgb(
+            rgb=fine_field_outputs.rgb, density=fine_field_outputs.density, deltas=pdf_ray_samples.deltas
+        )  # RendererOutputs
+        # outputs:
+        outputs = {"rgb_coarse": coarse_renderer_outputs.rgb, "rgb_fine": fine_renderer_outputs.rgb}
         return outputs
 
     def get_losses(self, batch, graph_outputs):
         # batch.pixels # (num_rays, 3)
         losses = {}
-        # rgb_loss_coarse = self.rgb_loss(batch.pixels, graph_outputs["rgb_coarse"])
-        # rgb_loss_fine = self.rgb_loss(batch.pixels, graph_outputs["rgb_fine"])
-        # losses = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine}
+        rgb_loss_coarse = self.rgb_loss(batch.pixels, graph_outputs["rgb_coarse"])
+        rgb_loss_fine = self.rgb_loss(batch.pixels, graph_outputs["rgb_fine"])
+        losses = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine}
         return losses

mattport/nerf/renderers.py

@@ -1,19 +1,28 @@
 """
 Collection of renderers
 """
-from typing import Dict, Set
+from dataclasses import dataclass
+from typing import Set
 
 import torch
 from torch import nn
 from torchtyping import TensorType
 
-from mattport.nerf.field_modules import field_outputs
+from mattport.nerf.field_modules import field_heads
+
+
+@dataclass
+class RendererOutputs:
+    """_summary_"""
+
+    rgb: TensorType["num_rays", 3] = None
+    density: TensorType["num_rays", 1] = None
 
 
 class Renderer(nn.Module):
     """Base Renderer. Intended to be subclassed"""
 
-    def required_field_outputs(self) -> Set[field_outputs.FieldOutput]:
+    def required_field_outputs(self) -> Set[field_heads.FieldHead]:
         """
         Returns:
             Set[RenderHead]: RenderHeads required for this renderer.
@@ -22,9 +31,10 @@ class Renderer(nn.Module):
 
     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"]:
+    ) -> RendererOutputs:
         """Composite samples along ray and render image
 
         Args:
@@ -40,18 +50,15 @@ class Renderer(nn.Module):
 class RGBRenderer(Renderer):
     """Standard volumetic rendering."""
 
-    def __init__(self) -> None:
-        super().__init__()
-
-    def required_field_outputs(self) -> Set[field_outputs.FieldOutput]:
-        return set(field_outputs.RGBFieldOutput, field_outputs.DensityFieldOutput)
+    def required_field_outputs(self) -> Set[field_heads.FieldHead]:
+        return set(field_heads.RGBFieldHead, field_heads.DensityFieldHead)
 
     def forward(
         self,
         rgb: TensorType[..., "num_samples", 3],
         density: TensorType[..., "num_samples", 1],
         deltas: TensorType[..., "num_samples"],
-    ) -> TensorType[..., "out_dim"]:
+    ) -> RendererOutputs:
         """Composite samples along ray and render color image
 
         Args:
@@ -65,14 +72,15 @@ class RGBRenderer(Renderer):
         alphas = 1 - torch.exp(-delta_density)
 
         transmittance = torch.cumsum(delta_density[..., :-1], dim=-1)
-        transmittance = torch.cat([torch.zeros((*transmittance.shape[:1], 1)), transmittance], axis=-1)
+        transmittance = torch.cat([torch.zeros((*transmittance.shape[:1], 1)).to(rgb.device), transmittance], axis=-1)
         transmittance = torch.exp(-transmittance)  # [..., "num_samples"]
 
         weights = alphas * transmittance  # [..., "num_samples"]
 
         rgb = torch.sum(weights[..., None] * rgb, dim=-2)
 
-        return rgb
+        renderer_outputs = RendererOutputs(rgb=rgb)
+        return renderer_outputs
 
 
 class DepthRenderer(Renderer):
@@ -82,12 +90,13 @@ class DepthRenderer(Renderer):
         Renderer (_type_): _description_
     """
 
-    def required_field_outputs(self) -> Set[field_outputs.FieldOutput]:
-        return set(field_outputs.DensityFieldOutput)
+    def required_field_outputs(self) -> Set[field_heads.FieldHead]:
+        return set(field_heads.DensityFieldHead)
 
     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"]:
+    ) -> RendererOutputs:
         raise NotImplementedError

mattport/nerf/sampler.py

@@ -1,7 +1,6 @@
 """
 Collection of sampling strategies
 """
-from typing import NamedTuple
 
 import torch
 from torch import nn
@@ -80,5 +79,6 @@ class UniformSampler(Sampler):
             t_min=torch.ones_like(ray_bundle.origins[:, 0]) * near_plane,
             t_max=torch.ones_like(ray_bundle.origins[:, 0]) * far_plane,
             ts=bins.unsqueeze(0).repeat(num_rays, 1),  # (num_rays, num_samples)
+            ray_bundle=ray_bundle
         )
         return ray_samples

mattport/nerf/trainer.py

@@ -1,9 +1,10 @@
 """
 Code to train model.
 """
-import pprint
 
+import logging
 import torch.distributed as dist
+from hydra.utils import instantiate
 from omegaconf import DictConfig
 from torch import optim
 from torch.nn.parallel import DistributedDataParallel as DDP
@@ -12,13 +13,13 @@ from tqdm import tqdm
 
 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 hydra.utils import instantiate
 
 
 class Trainer:
     """Training class"""
 
+    # pylint: disable=too-many-instance-attributes
+
     def __init__(self, config: DictConfig, local_rank: int = 0, world_size: int = 1):
         self.config = config
         self.local_rank = local_rank
@@ -88,7 +89,7 @@ class Trainer:
     def train_epoch(self):
         """_summary_"""
         num_iters = 100
-        for _ in range(num_iters):
+        for _ in tqdm(range(num_iters)):
             batch = next(iter(self.train_dataloader))
             self.train_iteration(batch)
 
@@ -99,10 +100,9 @@ class Trainer:
             pass
 
         # move batch to correct device
-        batch.indices = batch.indices.to(f"cuda:{self.local_rank}")
-
-        ray_indices = batch.indices
+        ray_indices = batch.indices.to(f"cuda:{self.local_rank}")
         graph_outputs = self.graph(ray_indices)
+        batch.pixels = batch.pixels.to(f"cuda:{self.local_rank}")
         losses = self.graph.get_losses(batch, graph_outputs)
-
-        # TODO(ethan): update the network parameters
+        logging.info(losses)
+        

mattport/structures/rays.py

@@ -5,7 +5,6 @@ import random
 from dataclasses import dataclass
 from typing import Optional
 
-import torch
 from torchtyping import TensorType
 
 
@@ -68,44 +67,31 @@ class CameraRayBundle:
     camera_index: int = None
 
 
-@dataclass
 class RaySamples:
     """_summary_
-
-    Args:
-        t_min: maximum times
-        t_max: minimum times
-        ts: time steps
     """
-
-    t_min: TensorType["num_rays"]
-    t_max: TensorType["num_rays"]
-    ts: TensorType["num_rays", "num_samples+1"]  # TODO(ethan): check with Matt about this. 2nd dimension is num_samples+1
+    def __init__(
+        self,
+        t_min: TensorType["num_rays"],
+        t_max: TensorType["num_rays"],
+        ts: TensorType["num_rays", "num_samples+1"],
+        ray_bundle: RayBundle,
+    ) -> None:
+        self.t_min = t_min
+        self.t_max = t_max
+        self.ts = ts
+        self.positions = self.get_positions(ray_bundle)
+        self.directions = ray_bundle.directions.unsqueeze(1).repeat(1, self.positions.shape[1], 1)
+        self.deltas = self.get_deltas()
 
     def get_positions(self, ray_bundle: RayBundle) -> TensorType["num_rays", "num_samples", 3]:
         """Returns positions."""
         # TODO(ethan): check with Matt about this
-        t_mid = (self.ts[:, 1:] + self.ts[:, :-1]) / 2 # (num_rays, num_samples)
+        t_mid = (self.ts[:, 1:] + self.ts[:, :-1]) / 2  # (num_rays, num_samples)
         return ray_bundle.origins[:, None] + t_mid[:, :, None] * ray_bundle.directions[:, None]
 
-    def get_starts(self):
-        """_summary_
-
-        Returns:
-            _type_: _description_
-        """
-        t_mid = (self.ts[:, 1:] + self.ts[:, :-1]) / 2
-        return torch.cat([self.t_min[:, None], t_mid], dim=-1)
-
-    def get_ends(self):
-        """_summary_
-
-        Returns:
-            _type_: _description_
-        """
-        t_mid = (self.ts[:, 1:] + self.ts[:, :-1]) / 2
-        return torch.cat([t_mid, self.t_max[:, None]], dim=-1)
-
     def get_deltas(self) -> TensorType["num_samples"]:
         """Returns deltas."""
-        return self.get_ends() - self.get_starts()
+        deltas = self.ts[..., 1:] - self.ts[..., :-1]
+        # TODO(ethan): check this with Matt
+        return deltas

tests/nerf/field_modules/test_field_outputs.py

@@ -5,7 +5,7 @@ import pytest
 import torch
 from torch import nn
 
-from mattport.nerf.field_modules import field_outputs
+from mattport.nerf.field_modules import field_heads
 
 
 def test_field_output():
@@ -13,7 +13,7 @@ def test_field_output():
     in_dim = 6
     out_dim = 4
     activation = nn.ReLU()
-    render_head = field_outputs.FieldOutput(in_dim=in_dim, out_dim=out_dim, activation=activation)
+    render_head = field_heads.FieldHead(in_dim=in_dim, out_dim=out_dim, activation=activation)
     assert render_head.get_out_dim() == out_dim
 
     x = torch.ones((9, in_dim))
@@ -29,7 +29,7 @@ def test_field_output():
 def test_density_head():
     """Test rgb output"""
     in_dim = 6
-    density_head = field_outputs.DensityFieldOutput(in_dim)
+    density_head = field_heads.DensityFieldHead(in_dim)
     assert density_head.get_out_dim() == 1
 
     x = torch.ones((9, in_dim))
@@ -44,7 +44,7 @@ def test_density_head():
 def test_rgb_head():
     """Test rgb output"""
     in_dim = 6
-    rgb_head = field_outputs.RGBFieldOutput(in_dim)
+    rgb_head = field_heads.RGBFieldHead(in_dim)
     assert rgb_head.get_out_dim() == 3
 
     x = torch.ones((9, in_dim))

tests/nerf/graph/test_graph.py

@@ -5,7 +5,7 @@ import numpy as np
 from hydra import compose, initialize
 from omegaconf import open_dict
 
-from mattport.nerf.graph import Graph, Node
+from mattport.nerf.graph.base import Graph, Node
 
 CONFIG_DIR = "configs"