Update Instant-NGP & Density Grid with CUDA supports. Much faster and better PSNR. (#193)

pyrad/fields/density_fields/base.py

@@ -0,0 +1,13 @@
+# Copyright 2022 The Plenoptix Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

pyrad/fields/density_fields/density_grid.py

@@ -0,0 +1,194 @@
+# Copyright 2022 The Plenoptix Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Code to implement the density grid.
+"""
+from typing import Callable, List, NoReturn, Tuple
+
+import torch
+from torch import nn
+from torchtyping import TensorType
+
+import pyrad.cuda as pyrad_cuda
+
+
+def _create_grid_coords(resolution: int, device: torch.device = "cpu") -> TensorType["n_coords", 3]:
+    """Create 3D grid coordinates
+
+    Args:
+        resolution (int): The 3D resolution of the grid.
+        device (torch.device): Which device you want the returned tensor to live int.
+
+    Returns:
+        TensorType["n_coords", 3]: All grid coordinates with shape [res * res * res, 3]
+    """
+    arrange = torch.arange(resolution, device=device)
+    grid_x, grid_y, grid_z = torch.meshgrid([arrange, arrange, arrange], indexing="ij")
+    coords = torch.stack([grid_x, grid_y, grid_z])  # [3, res, res, res]
+    coords = coords.reshape(3, -1).t().contiguous()  # [N, 3]
+    return coords
+
+
+class DensityGrid(nn.Module):
+    """Cascaded multi-res density grids.
+
+    The multi-res grids are all centerd at [center, center, center] in the world space.
+    The first grid covers regions (center - base_scale / 2, center + base_scale / 2)
+    and the followings grows by 2x scale each.
+
+    TODO(ruilongli): support `center` and `base_scale` with 3-dim.
+
+    Args:
+        center (float, optional): Center of all the grids in the world space. Defaults to 0.
+        base_scale (float, optional): Center of the scale of the base level grid. Defaults to 1.
+        num_cascades (int, optional): Number of the cascaded multi-res levels. Defaults to 1.
+        resolution (int, optional): Resolution of the grid. Defaults to 128.
+        update_every_num_iters (int, optional): How frequently to update the grid values. Defaults to 16.
+    """
+
+    def __init__(
+        self,
+        center: float = 0.0,
+        base_scale: float = 1.0,
+        num_cascades: int = 1,
+        resolution: int = 128,
+        update_every_num_iters: int = 16,
+    ) -> None:
+        super().__init__()
+        self.center = center
+        self.base_scale = base_scale
+        self.num_cascades = num_cascades  # the number of levels (i.e, cascades)
+        self.resolution = resolution
+        self.update_every_num_iters = update_every_num_iters
+
+        density_grid = torch.zeros([self.num_cascades] + [self.resolution**3])
+        self.register_buffer("density_grid", density_grid)
+
+        density_bitfield = torch.zeros([self.num_cascades * self.resolution**3 // 8], dtype=torch.uint8)
+        self.register_buffer("density_bitfield", density_bitfield)
+
+        # Integer grid coords / indices that do not related to cascades
+        grid_coords = _create_grid_coords(resolution)
+        # TODO(ruilongli): hacky way for now. support cpu version of morton3D?
+        grid_indices = pyrad_cuda.morton3D(grid_coords.to("cuda:0")).to("cpu")
+        self.register_buffer("grid_coords", grid_coords)
+        self.register_buffer("grid_indices", grid_indices)
+
+    @torch.no_grad()
+    def get_all_cells(self) -> List[Tuple[TensorType["n_cells"], TensorType["n_cells", 3]]]:
+        """Returns all cells of the grid.
+
+        Returns:
+            List[Tuple[TensorType["n_cells"], TensorType["n_cells", 3]]]: All grid indices and grid
+            coordinates of all cascade levels.
+        """
+        return [(self.grid_indices, self.grid_coords)] * self.num_cascades
+
+    @torch.no_grad()
+    def sample_uniform_and_occupied_cells(self, n: int) -> List[Tuple[TensorType["n"], TensorType["n", 3]]]:
+        """Samples both n uniform and occupied cells (per cascade)
+
+        Args:
+            n: Number of cells to be sampled.
+
+        Returns:
+            List[Tuple[TensorType["n"], TensorType["n", 3]]]: Sampled grid indices and grid
+            coordinates of all cascade levels.
+        """
+        device = self.density_grid.device
+
+        cells = []
+        for c in range(self.num_cascades):
+            uniform_coords = torch.randint(self.resolution, (n, 3), device=device)
+            uniform_indices = pyrad_cuda.morton3D(uniform_coords)
+
+            occupied_indices = torch.nonzero(self.density_grid[c] > 0)[:, 0]
+            if n < len(occupied_indices):
+                selector = torch.randint(len(occupied_indices), (n,), device=device)
+                occupied_indices = occupied_indices[selector]
+            occupied_coords = pyrad_cuda.morton3D_invert(occupied_indices)
+
+            # concatenate
+            coords = torch.cat([uniform_coords, occupied_coords], dim=0)
+            indices = torch.cat([uniform_indices, occupied_indices], dim=0)
+            cells.append((indices, coords))
+        return cells
+
+    @torch.no_grad()
+    def update_density_grid(
+        self,
+        density_eval_func: Callable,
+        step: int,
+        density_threshold: float = 2,  # 0.01 / (SQRT3 / 1024 * 3)
+        decay: float = 0.95,
+    ) -> NoReturn:
+        """Update the density grid in EMA way.
+
+        Args:
+            density_eval_func (Callable): A Callable function that takes in sample positions (N, 3) and
+                returns densities (N, 1).
+            step (int): The current training step. Instant-NGP has a warmup stage where all cells are updated.
+                After certain amount of steps (256), it speeds up by only update sampled cells.
+            density_threshold (float): The threshold to prune cells. Recommand to calculate it using this rule:
+                `weight_threshold / dt`. For example for Instant-NGP on Lego, the minimum step size `dt` is
+                `3` (scale of the scene) * `sqrt(3)` / `1024` (number of samples). And if we want to the weight
+                threshold to be `0.01`, we will get `0.01 / (3 * sqrt(3) / 1024) ~= 2.` for density threshold.
+            decay (float): EMA decay for density updating.
+        """
+        # create temporary grid
+        tmp_grid = -torch.ones_like(self.density_grid)
+        if step < 256:
+            cells = self.get_all_cells()
+        else:
+            N = self.resolution**3 // 4
+            cells = self.sample_uniform_and_occupied_cells(N)
+
+        # infer sigmas
+        for mip_level in range(self.num_cascades):
+            mip_scale = 2 ** (-mip_level) / self.base_scale
+            indices, coords = cells[mip_level]
+            # `coords` denotes the i-th cell. It's a poor naming here.
+            # the actually coordinates in world space x has mapping to
+            # the cell like this:
+            # x \in [0, 1/res] maps to 0-th cell.
+            # x \in [1 - 1/res, 1] maps to (res-1)-th cell.
+            x = (coords + torch.rand_like(coords.float())) / self.resolution
+            x = (x - 0.5) / mip_scale + self.center
+            tmp_grid[mip_level, indices] = density_eval_func(x).squeeze(-1)
+
+        # ema update
+        valid_mask = (self.density_grid >= 0) & (tmp_grid >= 0)
+        self.density_grid[valid_mask] = torch.maximum(self.density_grid[valid_mask] * decay, tmp_grid[valid_mask])
+        mean_density = self.density_grid.mean().item()
+
+        # pack to bitfield
+        self.density_bitfield.data = pyrad_cuda.packbits(self.density_grid, min(mean_density, density_threshold))
+
+        # TODO(ruilongli): max pooling? https://github.com/NVlabs/instant-ngp/blob/master/src/testbed_nerf.cu#L578
+
+    @torch.no_grad()
+    def get_densities(self, xyzs: TensorType[..., 3]) -> TensorType[..., 1]:
+        """Trilinear interpolation to get the density values.
+        Args:
+            xyzs (TensorType[..., 3]): 3D querry coordinate
+        Returns:
+            TensorType[..., 1]: Density values
+        """
+        # TODO(ruilongli) pass in RaySamples because we need dt to decide mip level.
+        raise NotImplementedError("We haven't implement the logic of querying densities from cascaded grid.")
+
+    def forward(self, x):
+        """Not implemented."""
+        raise RuntimeError("Shouldn't be called!")