nerfacc integration (#447)

* integrated changes

* update

* stratified

* black

* bump version

* upgrade version

.vscode/launch.json

@@ -20,7 +20,7 @@
       "justMyCode": true
     },
     {
-      "name": "Python: run_train (gpu)",
+      "name": "Python: run_train",
       "type": "python",
       "request": "launch",
       "program": "scripts/run_train.py",
@@ -30,17 +30,7 @@
       ]
     },
     {
-      "name": "Python: run_train (cpu)",
-      "type": "python",
-      "request": "launch",
-      "program": "scripts/run_train.py",
-      "console": "integratedTerminal",
-      "args": [
-        "machine.num_gpus=0"
-      ]
-    },
-    {
-      "name": "Python: run_train ngp (gpu)",
+      "name": "Python: run_train ngp",
       "type": "python",
       "request": "launch",
       "program": "scripts/run_train.py",
@@ -50,7 +40,7 @@
       ]
     },
     {
-      "name": "Python: run_train ngp nerfactory (gpu)",
+      "name": "Python: run_train ngp nerfactory",
       "type": "python",
       "request": "launch",
       "program": "scripts/run_train.py",
@@ -60,16 +50,16 @@
         "--experiment_name=nerfactory_experiment",
         "--logging.writer.2.max-log-size=0",
         // viewer
-        "--viewer.enable",
-        "--viewer.zmq-url=tcp://127.0.0.1:6001",
-        "--viewer.websocket-port=8014",
-        "--pipeline.model.randomize-background",
-        "--pipeline.model.cone-angle=0.00390625",
+        // "--viewer.enable",
+        // "--viewer.zmq-url=tcp://127.0.0.1:6001",
+        // "--viewer.websocket-port=8014",
+        // "--pipeline.model.randomize-background",
+        // "--pipeline.model.cone-angle=0.00390625",
         "nerfactory-data"
       ]
     },
     {
-      "name": "Python: run_train ngp-friends (gpu)",
+      "name": "Python: run_train ngp-friends",
       "type": "python",
       "request": "launch",
       "program": "scripts/run_train.py",
@@ -88,7 +78,7 @@
       ]
     },
     {
-      "name": "Python: run_train mipnerf 360 (gpu)",
+      "name": "Python: run_train mipnerf 360",
       "type": "python",
       "request": "launch",
       "program": "scripts/run_train.py",
@@ -98,7 +88,7 @@
       ]
     },
     {
-      "name": "Python: run_train tensorf (gpu)",
+      "name": "Python: run_train tensorf",
       "type": "python",
       "request": "launch",
       "program": "scripts/run_train.py",
@@ -108,7 +98,7 @@
       ]
     },
     {
-      "name": "Python: run_train nerfw (gpu)",
+      "name": "Python: run_train nerfw",
       "type": "python",
       "request": "launch",
       "program": "scripts/run_train.py",
@@ -118,7 +108,7 @@
       ]
     },
     {
-      "name": "Python: run_train semantic_nerf (gpu)",
+      "name": "Python: run_train semantic_nerf",
       "type": "python",
       "request": "launch",
       "program": "scripts/run_train.py",

nerfactory/configs/base.py

@@ -345,6 +345,10 @@ class VanillaDataManagerConfig(InstantiateConfig):
     """number of images to sample during training iteration"""
     eval_dataparser: dcargs.conf.Fixed[Optional[InstantiateConfig]] = None
     """optionally specify different dataparser to use during eval; if None, uses train_dataparser"""
+    eval_num_rays_per_batch: int = 1024
+    """number of rays per batch to use per eval iteration"""
+    eval_num_images_to_sample_from: int = -1
+    """number of images to sample during eval iteration"""
     eval_image_indices: Optional[Tuple[int, ...]] = (0,)
     """specifies the image indices to use during eval; if None, uses all"""
     eval_num_rays_per_chunk: int = 4096
@@ -407,8 +411,6 @@ class CompoundModelConfig(ModelConfig):
     field_implementation: Literal["torch", "tcnn"] = "tcnn"  # torch, tcnn, ...
     loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss": 1.0})
     num_samples: int = 1024  # instead of course/fine samples
-    cone_angle: float = 0.0
-    near_plane: float = 0.5
 
 
 @dataclass

nerfactory/datamanagers/base.py

@@ -235,7 +235,7 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         """Sets up the data loaders for training"""
         assert self.train_input_dataset is not None
         if self.world_size > 1:
-            self.sampler = DistributedSampler(
+            sampler = DistributedSampler(
                 self.train_input_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42
             )
             self.train_image_dataloader = CacheImageDataloader(
@@ -244,7 +244,7 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
                 device=self.device,
                 num_workers=self.world_size * 4,
                 pin_memory=True,
-                sampler=self.sampler,
+                sampler=sampler,
             )  # TODO(ethan): pass this in
         else:
             self.train_image_dataloader = CacheImageDataloader(
@@ -261,6 +261,30 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
     def setup_eval(self):
         """Sets up the data loader for evaluation"""
         assert self.eval_input_dataset is not None
+        if self.world_size > 1:
+            sampler = DistributedSampler(
+                self.eval_input_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42
+            )
+            self.eval_image_dataloader = CacheImageDataloader(
+                self.eval_input_dataset,
+                num_images_to_sample_from=self.config.eval_num_images_to_sample_from,
+                device=self.device,
+                num_workers=self.world_size * 4,
+                pin_memory=True,
+                sampler=sampler,
+            )  # TODO(ethan): pass this in
+        else:
+            self.eval_image_dataloader = CacheImageDataloader(
+                self.eval_input_dataset,
+                num_images_to_sample_from=self.config.eval_num_images_to_sample_from,
+                device=self.device,
+                num_workers=self.world_size * 4,
+                pin_memory=True,
+            )
+        self.iter_eval_image_dataloader = iter(self.eval_image_dataloader)
+        self.eval_pixel_sampler = PixelSampler(self.config.eval_num_rays_per_batch)
+        self.eval_ray_generator = RayGenerator(self.eval_input_dataset.dataset_inputs.cameras.to(self.device))
+        # for loading full images
         self.eval_dataloader = FixedIndicesEvalDataloader(
             input_dataset=self.eval_input_dataset,
             num_rays_per_chunk=self.config.eval_num_rays_per_chunk,
@@ -270,8 +294,7 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         )
 
     def next_train(self) -> Tuple[RayBundle, Dict]:
-        """Returns the next batch of data from the train dataloader.
-        The RayBundle can be shaped in whatever way."""
+        """Returns the next batch of data from the train dataloader."""
         self.train_count += 1
         image_batch = next(self.iter_train_image_dataloader)
         batch = self.train_pixel_sampler.sample(image_batch)
@@ -280,9 +303,10 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         return ray_bundle, batch
 
     def next_eval(self) -> Tuple[RayBundle, Dict]:
-        """Returns the next batch of data from the eval dataloader.
-        The RayBundle should be shaped like an image."""
+        """Returns the next batch of data from the eval dataloader."""
         self.eval_count += 1
-        assert self.eval_dataloader is not None, "Must setup eval dataloader before calling next_eval"
-        camera_ray_bundle, batch = next(self.eval_dataloader)
-        return camera_ray_bundle, batch
+        image_batch = next(self.iter_eval_image_dataloader)
+        batch = self.eval_pixel_sampler.sample(image_batch)
+        ray_indices = batch["indices"]
+        ray_bundle = self.eval_ray_generator(ray_indices)
+        return ray_bundle, batch

nerfactory/fields/instant_ngp_field.py

@@ -151,6 +151,26 @@ class TCNNInstantNGPField(Field):
         rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions)
         return {FieldHeadNames.RGB: rgb}
 
+    def get_outputs_from_positions_and_direction(self, positions: TensorType["bs", 3], directions: TensorType["bs", 3]):
+        """Given positions and directions, return the outputs of the field.
+
+        Args:
+            positions: positions of the points
+            directions: directions at the point locations
+        """
+        # assume flat positions and directions
+        # TODO: normalize positions
+        positions = SceneBounds.get_normalized_positions(positions, self.aabb)
+        directions = get_normalized_directions(directions)
+        h = self.mlp_base(positions)
+        density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1)
+        density = trunc_exp(density_before_activation.to(positions))
+        # TODO: normalize directions
+        d = self.direction_encoding(directions)
+        h = torch.cat([d, base_mlp_out], dim=-1)
+        rgb = self.mlp_head(h).to(directions)
+        return {FieldHeadNames.RGB: rgb, FieldHeadNames.DENSITY: density}
+
 
 class TorchInstantNGPField(NeRFField):
     """

nerfactory/models/instant_ngp.py

@@ -18,24 +18,25 @@ Implementation of Instant NGP.
 
 from __future__ import annotations
 
+import math
 from dataclasses import dataclass, field
 from typing import Dict, List, Literal, Optional, Type
 
+import nerfacc  # pylint: disable=import-error
 import torch
 from torch.nn import Parameter
 from torchmetrics import PeakSignalNoiseRatio
 from torchmetrics.functional import structural_similarity_index_measure
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from torchtyping import TensorType
 
-import nerfactory.cuda as nerfactory_cuda
 from nerfactory.cameras.rays import RayBundle
 from nerfactory.configs import base as cfg
 from nerfactory.fields.instant_ngp_field import field_implementation_to_class
 from nerfactory.fields.modules.field_heads import FieldHeadNames
 from nerfactory.models.base import Model
-from nerfactory.models.modules.ray_sampler import NGPSpacedSampler
 from nerfactory.optimizers.loss import MSELoss
-from nerfactory.utils import colors, visualization, writer
+from nerfactory.utils import visualization, writer
 from nerfactory.utils.callbacks import (
     TrainingCallback,
     TrainingCallbackAttributes,
@@ -57,7 +58,8 @@ class InstantNGPModelConfig(cfg.ModelConfig):
     """Instant NGP doesn't use a collider."""
     field_implementation: Literal["torch", "tcnn"] = "tcnn"  # torch, tcnn, ...
     """one of "torch" or "tcnn", or other fields in 'field_implementation_to_class'"""
-    enable_density_field: bool = True
+    # TODO(ethan): remove the density field specified here
+    enable_density_field: bool = False
     """Whether to create a density field to filter samples."""
     num_samples: int = 1024
     """Number of samples in field evaluation. Defaults to 1024,"""
@@ -83,27 +85,39 @@ class NGPModel(Model):
         self.field = None
         super().__init__(config=config, **kwargs)
 
-    def get_training_callbacks(
-        self, training_callback_attributes: TrainingCallbackAttributes
-    ) -> List[TrainingCallback]:
-        assert self.density_field is not None
-        return [
-            TrainingCallback(
-                where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
-                update_every_num_iters=self.density_field.update_every_num_iters,
-                func=self.density_field.update_density_grid,
-                kwargs={"density_eval_func": self.field.density_fn},  # type: ignore
-            )
-        ]
-
     def populate_modules(self):
         """Set the fields and modules."""
         super().populate_modules()
+
         # torch or tiny-cuda-nn version
         self.field = field_implementation_to_class[self.config.field_implementation](self.scene_bounds.aabb)
 
-        # samplers
-        self.sampler = NGPSpacedSampler(num_samples=self.config.num_samples, density_field=self.density_field)
+        # to match Ruilong's code naming
+        self.scene_aabb = Parameter(self.scene_bounds.aabb.flatten(), requires_grad=False)
+
+        # setup some rendering settings
+        render_n_samples = 1024
+        self.render_step_size = (
+            (self.scene_aabb[3:] - self.scene_aabb[:3]).max() * math.sqrt(3) / render_n_samples
+        ).item()
+
+        # setup occupancy field with eval function
+        def occ_eval_fn(x: torch.Tensor) -> torch.Tensor:
+            """Evaluate occupancy given positions.
+
+            Args:
+                x: positions with shape (N, 3).
+            Returns:
+                occupancy values with shape (N, 1).
+            """
+            density_after_activation = self.field.density_fn(x)
+            # those two are similar when density is small.
+            # occupancy = 1.0 - torch.exp(-density_after_activation * render_step_size)
+            occupancy = density_after_activation * self.render_step_size
+            return occupancy
+
+        # occupancy grid
+        self.occ_field = nerfacc.OccupancyField(occ_eval_fn=occ_eval_fn, aabb=self.scene_aabb, resolution=128)
 
         # losses
         self.rgb_loss = MSELoss()
@@ -113,6 +127,17 @@ class NGPModel(Model):
         self.ssim = structural_similarity_index_measure
         self.lpips = LearnedPerceptualImagePatchSimilarity()
 
+    def get_training_callbacks(
+        self, training_callback_attributes: TrainingCallbackAttributes
+    ) -> List[TrainingCallback]:
+        return [
+            TrainingCallback(
+                where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],
+                update_every_num_iters=1,
+                func=self.occ_field.every_n_step,  # will take in step
+            )
+        ]
+
     def get_param_groups(self) -> Dict[str, List[Parameter]]:
         param_groups = {}
         if self.field is None:
@@ -123,43 +148,34 @@ class NGPModel(Model):
     def get_outputs(self, ray_bundle: RayBundle):
         assert self.field is not None
 
-        # TODO(ruilongli)
-        # - train test difference
-        # - visualize "depth_density_grid"
-        num_rays = len(ray_bundle)
-        device = ray_bundle.origins.device
+        def query_fn(positions: TensorType["bs", 3], directions: TensorType["bs", 3], only_density=False):
+            if only_density:
+                return self.field.density_fn(positions)
+            field_outputs = self.field.get_outputs_from_positions_and_direction(positions, directions)
+            rgbs = field_outputs[FieldHeadNames.RGB]
+            sigmas = field_outputs[FieldHeadNames.DENSITY]
+            return rgbs, sigmas
 
-        ray_samples, packed_info, t_min, t_max = self.sampler(
-            ray_bundle, self.field.aabb, cone_angle=self.config.cone_angle, near_plane=self.config.near_plane
-        )
-
-        field_outputs = self.field.forward(ray_samples)
-        rgbs = field_outputs[FieldHeadNames.RGB]
-        sigmas = field_outputs[FieldHeadNames.DENSITY]
-
-        # accumulate all the rays start from zero opacity
-        opacities = torch.zeros((num_rays, 1), device=device)
         (
-            accumulated_weight,
-            accumulated_depth,
             accumulated_color,
-            alive_ray_mask,
-        ) = nerfactory_cuda.VolumeRenderer.apply(
-            packed_info,
-            ray_samples.frustums.starts,
-            ray_samples.frustums.ends,
-            sigmas.contiguous(),
-            rgbs.contiguous(),
-            opacities,
+            accumulated_depth,
+            accumulated_weight,
+            steps_counter,  # pylint: disable=unused-variable
+            compact_steps_counter,  # pylint: disable=unused-variable
+        ) = nerfacc.volumetric_rendering(
+            query_fn=query_fn,
+            rays_o=ray_bundle.origins,
+            rays_d=ray_bundle.directions,
+            scene_aabb=self.scene_aabb,
+            scene_occ_binary=self.occ_field.occ_grid_binary,
+            scene_resolution=self.occ_field.resolution,
+            render_bkgd=torch.ones(3, device=self.device),
+            render_step_size=self.render_step_size,
+            near_plane=self.config.near_plane,
+            stratified=self.training,  # only use stratified sampling during training
         )
-        accumulated_depth = torch.clip(accumulated_depth, t_min[:, None], t_max[:, None])
 
-        if self.config.randomize_background:
-            background_colors = torch.rand((num_rays, 3), device=device)  # (num_rays, 3)
-        else:
-            background_colors = colors.WHITE.to(accumulated_color)  # (3,)
-
-        accumulated_color = accumulated_color + background_colors * (1.0 - accumulated_weight)
+        alive_ray_mask = accumulated_weight.squeeze(-1) > 0
 
         outputs = {
             "rgb": accumulated_color,

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "nerfactory"
-version = "0.0.0"
+version = "0.0.1"
 description = "All-in-one repository for state-of-the-art NeRFs"
 readme = "README.md"
 license = { text="Apache 2.0"}
@@ -33,7 +33,7 @@ dependencies = [
     "mediapy==1.1.0",
     "msgpack==1.0.4",
     "msgpack_numpy==0.4.8",
-    "nerfacc==0.0.3",
+    "nerfacc==0.0.5",
     "numpy-quaternion",
     "opencv-python==4.6.0.66",
     "plotly==5.7.0",