Merge pull request #51 from ethanweber/tancik/ngp_encoding

Add basic NGP encoder

.github/workflows/code_checks.yml

@@ -37,7 +37,8 @@ jobs:
             --disable=too-many-arguments \
             --disable=too-many-locals \
             --disable=fixme \
-            --disable=too-many-instance-attributes
+            --disable=too-many-instance-attributes \
+            --disable=duplicate-code
       - name: Test with pytest
         run: |
           pytest tests --typeguard-packages=mattport --torchtyping-patch-typeguard --disable-warnings

.vscode/launch.json

@@ -3,7 +3,8 @@
     // Hover to view descriptions of existing attributes.
     // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
     "version": "0.2.0",
-    "configurations": [{
+    "configurations": [
+        {
             "name": "Python: Current File",
             "type": "python",
             "request": "launch",
@@ -27,6 +28,16 @@
             "args": [
                 "machine_config.num_gpus=0"
             ]
-        }
+        },
+        {
+            "name": "Python: run_train ngp (gpu)",
+            "type": "python",
+            "request": "launch",
+            "program": "scripts/run_train.py",
+            "console": "integratedTerminal",
+            "args": [
+                "--config-name=instant_ngp.yaml"
+            ]
+        },
     ]
 }

configs/graph/instant_ngp_graph.yaml

@@ -0,0 +1,29 @@
+# saving model specifications
+model_dir: /tmp/mattport_models/ 
+steps_per_save: null
+steps_per_test: 100
+max_num_iterations: 1000000
+
+# pretrained loading options
+resume_train:
+  load_dir: # pretrained model directory
+  load_step: # pretrained model's step number
+
+# NOTE(ethan): currently setting to nothing b/c we're still designing the config format...
+network:
+  _target_: mattport.nerf.graph.instant_ngp.NGPGraph
+  near_plane: 1.0
+  far_plane: 6.0
+  num_coarse_samples: 64
+  num_importance_samples: 64
+
+# optimizer options for network
+param_groups:
+  fields:
+    optimizer:
+      type: Adam
+      lr: 1e-2
+    scheduler:
+      type: StepLR
+      step_size: 30 #  Period of learning rate decay.
+      gamma: 0.1 # Multiplicative factor of learning rate decay. Default: 0.1.

configs/instant_ngp.yaml

@@ -0,0 +1,23 @@
+defaults:
+  # This makes the logging module logs colorful
+  - data: default_data.yaml
+  - graph: instant_ngp_graph.yaml
+  - logging: default_logging.yaml
+  - override hydra/job_logging: colorlog
+  - override hydra/hydra_logging: colorlog
+
+hydra:
+  run:
+    # Configure output dir of each experiment programmatically from the arguments
+    # Example "outputs/mnist/classifier/baseline/2021-03-10-141516"
+    dir: outputs/${experiment_name}/${method_name}/${now:%Y-%m-%d_%H%M%S}
+
+experiment_name: blender_lego
+method_name: instant_ngp
+
+# machine options
+machine_config:
+  num_gpus: 1
+  num_machines: 1
+  machine_rank: 0
+  dist_url: auto

mattport/nerf/field_modules/encoding.py

@@ -1,6 +1,7 @@
 """
 Encoding functions
 """
+import numpy as np
 import torch
 from torchtyping import TensorType
 from mattport.nerf.field_modules.base import FieldModule
@@ -141,13 +142,96 @@ class RFFEncoding(Encoding):
 class HashEncoding(Encoding):
     """Hash encoding"""
 
-    def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
-        """Encodes an input tensor.
+    def __init__(
+        self,
+        num_levels: int = 16,
+        min_res: int = 16,
+        max_res: int = 1024,
+        hash_table_size: int = 2**19,
+        features_per_level: int = 2,
+        hash_init_scale: float = 0.001,
+    ) -> None:
+        """_summary_
 
         Args:
-            x (TensorType[..., "input_dim"]): Input tensor to be encoded
+            num_levels (int, optional): Number of feature grids. Defaults to 16.
+            min_res (int, optional): Resolution of smallest feature grid. Defaults to 16.
+            max_res (int, optional): Resolution of largest feature grid. Defaults to 1024.
+            hash_table_size (int, optional): Size of hash table. Defaults to 2**19.
+            features_per_level (int, optional): Number of features per level. Defaults to 2.
+            hash_init_scale (float, optional): Value to initialize hash grid. Defaults to 0.001.
+        """
+        super().__init__(in_dim=3)
+        self.num_levels = num_levels
+        self.features_per_level = features_per_level
+        self.hash_table_size = hash_table_size
+
+        levels = torch.arange(num_levels)
+        growth_factor = np.exp((np.log(max_res) - np.log(min_res)) / (num_levels - 1))
+        self.scalings = torch.floor(min_res * growth_factor**levels)
+
+        self.hash_offset = levels * hash_table_size
+        self.hash_table = torch.rand(size=(hash_table_size * num_levels, features_per_level)) * 2 - 1
+        self.hash_table *= hash_init_scale
+        self.hash_table = torch.nn.Parameter(self.hash_table)
+
+    def get_out_dim(self) -> int:
+        return self.num_levels * self.features_per_level
+
+    def hash_fn(self, in_tensor: TensorType[..., "num_levels", 3]) -> TensorType[..., "num_levels"]:
+        """Hash tensor using method described in Instant-NGP
+
+        Args:
+            in_tensor (TensorType[..., "num_levels", 3]): Tensor to be hashed
 
         Returns:
-            TensorType[..., "output_dim"]: A encoded input tensor
+            TensorType[..., "num_levels"]: Hashed tensor
         """
-        raise NotImplementedError
+        in_tensor = in_tensor * torch.tensor([1, 2654435761, 805459861]).to(in_tensor.device)
+        x = torch.bitwise_xor(in_tensor[..., 0], in_tensor[..., 1])
+        x = torch.bitwise_xor(x, in_tensor[..., 2])
+        x %= self.hash_table_size
+        x += self.hash_offset.to(x.device)
+        return x
+
+    def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
+        assert in_tensor.shape[-1] == 3
+
+        in_tensor = in_tensor[..., None, :]  # [..., 1, 3]
+        scaled = in_tensor * self.scalings.view(1, 1, -1, 1).to(in_tensor.device)  # [..., L, 3]
+        scaled_c = torch.ceil(scaled).type(torch.int32)
+        scaled_f = torch.floor(scaled).type(torch.int32)
+
+        offset = scaled - scaled_f
+
+        hashed_0 = self.hash_fn(scaled_c)  # [..., num_levels]
+        hashed_1 = self.hash_fn(torch.cat([scaled_c[..., 0:1], scaled_f[..., 1:2], scaled_c[..., 2:3]], dim=-1))
+        hashed_2 = self.hash_fn(torch.cat([scaled_f[..., 0:1], scaled_f[..., 1:2], scaled_c[..., 2:3]], dim=-1))
+        hashed_3 = self.hash_fn(torch.cat([scaled_f[..., 0:1], scaled_c[..., 1:2], scaled_c[..., 2:3]], dim=-1))
+        hashed_4 = self.hash_fn(torch.cat([scaled_c[..., 0:1], scaled_c[..., 1:2], scaled_f[..., 2:3]], dim=-1))
+        hashed_5 = self.hash_fn(torch.cat([scaled_c[..., 0:1], scaled_f[..., 1:2], scaled_f[..., 2:3]], dim=-1))
+        hashed_6 = self.hash_fn(scaled_f)
+        hashed_7 = self.hash_fn(torch.cat([scaled_f[..., 0:1], scaled_c[..., 1:2], scaled_f[..., 2:3]], dim=-1))
+
+        f_0 = self.hash_table[hashed_0]  # [..., num_levels, features_per_level]
+        f_1 = self.hash_table[hashed_1]
+        f_2 = self.hash_table[hashed_2]
+        f_3 = self.hash_table[hashed_3]
+        f_4 = self.hash_table[hashed_4]
+        f_5 = self.hash_table[hashed_5]
+        f_6 = self.hash_table[hashed_6]
+        f_7 = self.hash_table[hashed_7]
+
+        f_03 = f_0 * offset[..., 0:1] + f_3 * (1 - offset[..., 0:1])
+        f_12 = f_1 * offset[..., 0:1] + f_2 * (1 - offset[..., 0:1])
+        f_56 = f_5 * offset[..., 0:1] + f_6 * (1 - offset[..., 0:1])
+        f_47 = f_4 * offset[..., 0:1] + f_7 * (1 - offset[..., 0:1])
+
+        f0312 = f_03 * offset[..., 1:2] + f_12 * (1 - offset[..., 1:2])
+        f4756 = f_47 * offset[..., 1:2] + f_56 * (1 - offset[..., 1:2])
+
+        encoded_value = f0312 * offset[..., 2:3] + f4756 * (
+            1 - offset[..., 2:3]
+        )  # [..., num_levels, features_per_level]
+
+        return torch.flatten(encoded_value, start_dim=-2, end_dim=-1)  # [..., num_levels * features_per_level]

mattport/nerf/graph/instant_ngp.py

@@ -1,4 +1,161 @@
 """
-Instant NGP implementation.
-TODO:
+Implementation of Instant NGP.
 """
+
+
+from typing import Dict, List
+
+import torch
+from torch import nn
+from torch.nn import Parameter
+from torchtyping import TensorType
+
+from mattport.nerf.field_modules.encoding import HashEncoding, NeRFEncoding
+from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHeadNames, RGBFieldHead
+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
+from mattport.nerf.renderers import RGBRenderer
+from mattport.nerf.sampler import PDFSampler, UniformSampler  # pylint: disable=unused-import
+from mattport.structures.rays import RaySamples
+
+
+class NGPField(nn.Module):
+    """NeRF module"""
+
+    def __init__(self, num_layers=3, layer_width=64) -> None:
+        super().__init__()
+        self.num_layers = num_layers
+        self.layer_width = layer_width
+
+        self.encoding_xyz = HashEncoding()
+        self.encoding_dir = NeRFEncoding(
+            in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True
+        )
+        self.mlp_base = MLP(
+            in_dim=self.encoding_xyz.get_out_dim(),
+            out_dim=self.layer_width,
+            num_layers=self.num_layers,
+            layer_width=self.layer_width,
+            activation=nn.ReLU(),
+        )
+        self.mlp_rgb = MLP(
+            in_dim=self.mlp_base.get_out_dim() + self.encoding_dir.get_out_dim(),
+            out_dim=self.layer_width // 2,
+            num_layers=1,
+            activation=nn.ReLU(),
+        )
+        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
+        """
+        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)
+
+        field_outputs = {}
+        field_outputs.update(field_rgb_output)
+        field_outputs.update(field_density_out)
+        return field_outputs
+
+
+class NGPGraph(Graph):
+    """Vanilla NeRF graph"""
+
+    def __init__(
+        self,
+        intrinsics=None,
+        camera_to_world=None,
+        near_plane=1.0,
+        far_plane=6.0,
+        num_coarse_samples=64,
+        num_importance_samples=128,
+        **kwargs
+    ) -> None:
+        self.near_plane = near_plane
+        self.far_plane = far_plane
+        self.num_coarse_samples = num_coarse_samples
+        self.num_importance_samples = num_importance_samples
+        super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world, **kwargs)
+
+    def populate_modules(self):
+        # ray generator
+        self.ray_generator = RayGenerator(self.intrinsics, self.camera_to_world)
+
+        # samplers
+        self.sampler_uniform = UniformSampler(
+            near_plane=self.near_plane, far_plane=self.far_plane, num_samples=self.num_coarse_samples
+        )
+        self.sampler_pdf = PDFSampler(num_samples=self.num_importance_samples)
+
+        # field
+        self.field_coarse = NGPField()
+        self.field_fine = NGPField()
+
+        # renderers
+        self.renderer_rgb = RGBRenderer()
+
+        # losses
+        self.rgb_loss = MSELoss()
+
+    def get_param_groups(self) -> Dict[str, List[Parameter]]:
+        """Obtain the parameter groups for the optimizers
+
+        Returns:
+            Dict[str, List[Parameter]]: Mapping of different parameter groups
+        """
+        param_groups = {}
+        param_groups["fields"] = list(self.field_coarse.parameters()) + list(self.field_fine.parameters())
+        return param_groups
+
+    def forward(self, ray_indices: TensorType["num_rays", 3]):
+        """Takes in the ray indices and renders out values."""
+        # get the rays:
+        ray_bundle = self.ray_generator.forward(ray_indices)  # RayBundle
+        # coarse network:
+        uniform_ray_samples = self.sampler_uniform(ray_bundle)  # RaySamples
+
+        # time coarse network # xyz -> asd
+        coarse_field_outputs = self.field_coarse(uniform_ray_samples)  # FieldOutputs
+        # time end coarse network
+
+        coarse_weights = uniform_ray_samples.get_weights(coarse_field_outputs[FieldHeadNames.DENSITY])
+
+        coarse_renderer_outputs = self.renderer_rgb(
+            rgb=coarse_field_outputs[FieldHeadNames.RGB],
+            weights=coarse_weights,
+        )  # RendererOutputs
+        # fine network:
+        pdf_ray_samples = self.sampler_pdf(uniform_ray_samples, coarse_weights)  # RaySamples
+        fine_field_outputs = self.field_fine(pdf_ray_samples)  # FieldOutputs
+
+        fine_weights = pdf_ray_samples.get_weights(fine_field_outputs[FieldHeadNames.DENSITY])
+
+        fine_renderer_outputs = self.renderer_rgb(
+            rgb=fine_field_outputs[FieldHeadNames.RGB],
+            weights=fine_weights,
+        )  # 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}
+        return losses