revert to the ray-batched inference (#219)

configs/graph_instant_ngp.yaml

@@ -8,7 +8,7 @@ data:
     pixel_sampler:
       num_rays_per_batch: 8192 # for training
   dataloader_eval:
-    num_rays_per_chunk: 1000000000 # all rays for evaluation
+    num_rays_per_chunk: 8192 # all rays for evaluation
 
 viewer:
   enable: true

nerfactory/graphs/instant_ngp.py

@@ -86,82 +86,24 @@ class NGPGraph(Graph):
         num_rays = len(ray_bundle)
         device = ray_bundle.origins.device
 
-        if self.training:
-            ray_samples, packed_info, t_min, t_max = self.sampler(ray_bundle, self.field.aabb)
+        ray_samples, packed_info, t_min, t_max = self.sampler(ray_bundle, self.field.aabb)
 
-            field_outputs = self.field.forward(ray_samples)
-            rgbs = field_outputs[FieldHeadNames.RGB]
-            sigmas = field_outputs[FieldHeadNames.DENSITY]
+        field_outputs = self.field.forward(ray_samples)
+        rgbs = field_outputs[FieldHeadNames.RGB]
+        sigmas = field_outputs[FieldHeadNames.DENSITY]
 
-            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, rgbs, opacities
-            )
-            accumulated_depth = torch.clip(accumulated_depth, t_min[:, None], t_max[:, None])
-            accumulated_color = accumulated_color + colors.WHITE.to(accumulated_color) * (1.0 - accumulated_weight)
-
-        else:
-            accumulated_weight = torch.zeros((num_rays, 1), device=device)
-            accumulated_depth = torch.zeros((num_rays, 1), device=device)
-            accumulated_color = torch.zeros((num_rays, 3), device=device)
-            alive_ray_mask = torch.ones((num_rays), device=device, dtype=torch.bool)  # the ray we kept from sampler
-
-            t_min, t_max = None, None
-            t_marching = None
-            samples = 0
-            while samples < self.sampler.num_samples:
-                num_alive = alive_ray_mask.sum().item()
-                if num_alive == 0:
-                    break
-
-                min_samples = 1
-                num_samples = max(min(num_rays // num_alive, 64), min_samples)
-                samples += num_samples
-
-                ray_samples, packed_info, t_marching, t_max = self.sampler(
-                    ray_bundle, self.field.aabb, marching_steps=num_samples, t_min=t_marching, t_max=t_max
-                )
-                if t_min is None:
-                    t_min = t_marching
-
-                field_outputs = self.field.forward(ray_samples)
-                rgbs = field_outputs[FieldHeadNames.RGB]
-                sigmas = field_outputs[FieldHeadNames.DENSITY]
-
-                (
-                    _accumulated_weight,
-                    _accumulated_depth,
-                    _accumulated_color,
-                    _,
-                ) = nerfactory_cuda.VolumeRenderer.apply(
-                    packed_info,
-                    ray_samples.frustums.starts,
-                    ray_samples.frustums.ends,
-                    sigmas,
-                    rgbs,
-                    accumulated_weight,
-                )
-                accumulated_weight += _accumulated_weight
-                accumulated_depth += _accumulated_depth
-                accumulated_color += _accumulated_color
-
-                # all samples for this ray have been drawn
-                alive_ray_mask[packed_info[:, 2] < num_samples] = False
-
-                # this ray has converged
-                alive_ray_mask[accumulated_weight[:, 0] > 1.0 - 1e-4] = False
-
-                # march forward and skip those are not alive
-                t_marching = ray_samples.frustums.ends.reshape(num_rays, num_samples).max(dim=-1).values
-                t_marching[~alive_ray_mask] = 1e10
-
-            accumulated_depth = torch.clip(accumulated_depth, t_min[:, None], t_max[:, None])
-            accumulated_color = accumulated_color + colors.WHITE.to(accumulated_color) * (1.0 - accumulated_weight)
+        # 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, rgbs, opacities
+        )
+        accumulated_depth = torch.clip(accumulated_depth, t_min[:, None], t_max[:, None])
+        accumulated_color = accumulated_color + colors.WHITE.to(accumulated_color) * (1.0 - accumulated_weight)
 
         outputs = {
             "rgb": accumulated_color,

nerfactory/graphs/modules/ray_sampler.py

@@ -390,14 +390,10 @@ class NGPSpacedSampler(Sampler):
             t_min = torch.clamp(t_min, max=1e10)
             t_max = torch.clamp(t_max, max=1e10)
 
-        # sampling arguments
-        if self.training:
-            # TODO(ruilongli): * 16 is for original impl for training. Need to run
-            # some profiling test with this choice.
-            marching_steps = -1
-            max_samples_per_batch = len(rays_o) * num_samples
-        else:
-            max_samples_per_batch = len(rays_o) * marching_steps
+        marching_steps = -1  # disable marching mode for now
+        # TODO(ruilongli): * 16 is for original impl for training. Need to run
+        # some profiling test with this choice.
+        max_samples_per_batch = len(rays_o) * num_samples
 
         # marching
         packed_info, origins, dirs, starts, ends = nerfactory_cuda.raymarching(
@@ -421,20 +417,15 @@ class NGPSpacedSampler(Sampler):
         )
 
         # squeeze valid samples
-        if self.training:
-            total_samples = max(packed_info[:, -1].sum(), 1)
-            origins = origins[:total_samples]
-            dirs = dirs[:total_samples]
-            starts = starts[:total_samples]
-            ends = ends[:total_samples]
-            valid_mask = None
-        else:
-            valid_mask = (dirs != 0.0).any(dim=-1, keepdim=True)
+        total_samples = max(packed_info[:, -1].sum(), 1)
+        origins = origins[:total_samples]
+        dirs = dirs[:total_samples]
+        starts = starts[:total_samples]
+        ends = ends[:total_samples]
 
         # return samples
         zeros = torch.zeros_like(origins[:, :1])
         ray_samples = RaySamples(
             frustums=Frustums(origins=origins, directions=dirs, starts=starts, ends=ends, pixel_area=zeros),
-            valid_mask=valid_mask,
         )
         return ray_samples, packed_info, t_min, t_max