From c27dc26e036dc9c63c69ba807c011d80ab511218 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ruilong=20Li=28=E6=9D=8E=E7=91=9E=E9=BE=99=29?= Date: Wed, 27 Jul 2022 15:07:54 -0700 Subject: [PATCH] revert to the ray-batched inference (#219) --- configs/graph_instant_ngp.yaml | 2 +- nerfactory/graphs/instant_ngp.py | 90 +++++------------------- nerfactory/graphs/modules/ray_sampler.py | 27 +++---- 3 files changed, 26 insertions(+), 93 deletions(-) diff --git a/configs/graph_instant_ngp.yaml b/configs/graph_instant_ngp.yaml index 30e792a50..91ed03068 100644 --- a/configs/graph_instant_ngp.yaml +++ b/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 diff --git a/nerfactory/graphs/instant_ngp.py b/nerfactory/graphs/instant_ngp.py index 12eba5743..878033b09 100644 --- a/nerfactory/graphs/instant_ngp.py +++ b/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, diff --git a/nerfactory/graphs/modules/ray_sampler.py b/nerfactory/graphs/modules/ray_sampler.py index 1e464359b..6ac0c9efc 100644 --- a/nerfactory/graphs/modules/ray_sampler.py +++ b/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